From 36939cd53ac3c2aeaf943c7a14aa25781056ceab Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 09:38:55 +0800 Subject: [PATCH 01/10] Enhance Search Execution Results with DSL Handling - Added a `Results` field to the `SearchExecutionResult` struct to store raw search results, facilitating the extraction of DSL from database searches. - Implemented logic in the `saveSearch` method to extract and store the first DSL from DB search results, improving data retention for search operations. - Introduced a `dslToMap` method in the DB handler to convert QueryDSL to a map format for storage, enhancing the flexibility of result handling. - Updated the `Result` struct to include a `DSL` field for generated QueryDSL, ensuring comprehensive data representation for database search results. --- agent/assistant/search.go | 12 ++++++++++++ agent/search/handlers/db/handler.go | 24 ++++++++++++++++++++++++ agent/search/types/types.go | 3 +++ 3 files changed, 39 insertions(+) diff --git a/agent/assistant/search.go b/agent/assistant/search.go index c0bc05a2..4e344225 100644 --- a/agent/assistant/search.go +++ b/agent/assistant/search.go @@ -518,6 +518,7 @@ func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context Keywords: extractedKeywords, Config: ast.configToMap(searchConfig), RefCtx: refCtx, + Results: results, Duration: duration, SearchType: "auto", }) @@ -925,6 +926,7 @@ type SearchExecutionResult struct { Keywords []searchTypes.Keyword // Extracted keywords with weights Config map[string]any // Search config used RefCtx *searchTypes.ReferenceContext // Reference context with results + Results []*searchTypes.Result // Raw search results (for extracting DSL, etc.) Duration int64 // Search duration in ms Error error // Error if failed SearchType string // "auto", "web", "kb", "db" @@ -1012,6 +1014,16 @@ func (ast *Assistant) saveSearch(ctx *context.Context, execResult *SearchExecuti searchRecord.Prompt = execResult.RefCtx.Prompt } + // Extract DSL from DB search results + if execResult.Results != nil { + for _, result := range execResult.Results { + if result != nil && result.Type == searchTypes.SearchTypeDB && result.DSL != nil { + searchRecord.DSL = result.DSL + break // Only store the first DSL (usually there's only one DB search) + } + } + } + // Save to store if err := store.SaveSearch(searchRecord); err != nil { ctx.Logger.Warn("Failed to save search record: %v", err) diff --git a/agent/search/handlers/db/handler.go b/agent/search/handlers/db/handler.go index a728d39e..3bb8bf6c 100644 --- a/agent/search/handlers/db/handler.go +++ b/agent/search/handlers/db/handler.go @@ -207,6 +207,9 @@ func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Reques items = items[:maxResults] } + // 7. Convert DSL to map for storage + dslMap := h.dslToMap(result.DSL) + return &types.Result{ Type: types.SearchTypeDB, Query: req.Query, @@ -214,6 +217,7 @@ func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Reques Items: items, Total: len(items), Duration: time.Since(start).Milliseconds(), + DSL: dslMap, }, nil } @@ -377,3 +381,23 @@ func (h *Handler) extractContent(rec map[string]interface{}, mod *model.Model) s } return string(content) } + +// dslToMap converts QueryDSL to map for storage +func (h *Handler) dslToMap(dsl *gou.QueryDSL) map[string]interface{} { + if dsl == nil { + return nil + } + + // Marshal and unmarshal to get a clean map + data, err := json.Marshal(dsl) + if err != nil { + return nil + } + + var result map[string]interface{} + if err := json.Unmarshal(data, &result); err != nil { + return nil + } + + return result +} diff --git a/agent/search/types/types.go b/agent/search/types/types.go index 6cb70207..2b8237de 100644 --- a/agent/search/types/types.go +++ b/agent/search/types/types.go @@ -81,6 +81,9 @@ type Result struct { // Graph associations (KB only, if enabled) GraphNodes []*GraphNode `json:"graph_nodes,omitempty"` + + // DB specific + DSL map[string]interface{} `json:"dsl,omitempty"` // Generated QueryDSL (DB only) } // ResultItem represents a single search result item From 01820d9fe223b2894e1cbe631b9b537610ac7db8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 14:50:53 +0800 Subject: [PATCH 02/10] Add function execution support to Job system - Introduced `AddFunc` method to the `Job` struct for adding Go functions as job executions, allowing for dynamic execution of functions with specified arguments. - Enhanced internal execution handling to register functions in a global registry, ensuring proper cleanup after execution. - Implemented `ExecuteFunc` method in the `Goroutine` struct to handle the execution of registered functions, including error handling and context management. - Added comprehensive unit tests for `AddFunc`, verifying function registration, execution, and memory cleanup post-execution. - Updated related documentation to reflect the new functionality and usage patterns for adding and executing Go functions within the job system. --- job/execution.go | 29 +- job/goroutine.go | 78 +++++ job/job.go | 4 +- job/job_test.go | 271 ++++++++++++++++++ job/types.go | 55 +++- job/worker.go | 12 +- kb/api/addfile.go | 303 ++++++++++++++++++++ kb/api/addfile_test.go | 585 ++++++++++++++++++++++++++++++++++++++ kb/api/addtext.go | 230 +++++++++++++++ kb/api/addtext_test.go | 485 +++++++++++++++++++++++++++++++ kb/api/addurl.go | 230 +++++++++++++++ kb/api/addurl_test.go | 499 ++++++++++++++++++++++++++++++++ kb/api/collection_test.go | 9 +- kb/api/consts.go | 48 ++++ kb/api/document.go | 279 ++++++++++++++++++ kb/api/document_test.go | 291 +++++++++++++++++++ kb/api/interfaces.go | 21 +- kb/api/types.go | 124 ++++++++ kb/api/utils.go | 325 +++++++++++++++++++++ openapi/kb/addfile.go | 573 ++++++++++++------------------------- openapi/kb/addtext.go | 499 ++++++++++++-------------------- openapi/kb/addurl.go | 499 ++++++++++++-------------------- openapi/kb/document.go | 338 +++++----------------- openapi/kb/utils.go | 127 --------- 24 files changed, 4463 insertions(+), 1451 deletions(-) create mode 100644 kb/api/addfile.go create mode 100644 kb/api/addfile_test.go create mode 100644 kb/api/addtext.go create mode 100644 kb/api/addtext_test.go create mode 100644 kb/api/addurl.go create mode 100644 kb/api/addurl_test.go create mode 100644 kb/api/document.go create mode 100644 kb/api/document_test.go create mode 100644 kb/api/utils.go diff --git a/job/execution.go b/job/execution.go index fd75c488..bd33d261 100644 --- a/job/execution.go +++ b/job/execution.go @@ -29,6 +29,19 @@ func (j *Job) AddCommand(options *ExecutionOptions, command string, args []strin }) } +// AddFunc adds a new execution with a Go function +// The function is registered in a global registry and will be cleaned up after execution +// Note: The function is stored in memory registry and will be lost if the process restarts +func (j *Job) AddFunc(options *ExecutionOptions, name string, fn ExecutionFunc, args map[string]interface{}) error { + // fn will be registered in addExecution after ExecutionID is generated + return j.addExecution(options, &ExecutionConfig{ + Type: ExecutionTypeFunc, + Func: fn, // Temporarily store here, will be moved to registry + FuncName: name, + FuncArgs: args, + }) +} + // addExecution is the internal method to create execution records func (j *Job) addExecution(options *ExecutionOptions, config *ExecutionConfig) error { // Set default options if nil @@ -39,6 +52,14 @@ func (j *Job) addExecution(options *ExecutionOptions, config *ExecutionConfig) e } } + // For ExecutionTypeFunc, we need to register the function after getting ExecutionID + // Store the function temporarily and clear it before serialization + var funcToRegister ExecutionFunc + if config.Type == ExecutionTypeFunc && config.Func != nil { + funcToRegister = config.Func + config.Func = nil // Clear before serialization (can't be serialized anyway) + } + // Serialize ExecutionConfig to JSON for ConfigSnapshot configBytes, err := jsoniter.Marshal(config) if err != nil { @@ -61,11 +82,17 @@ func (j *Job) addExecution(options *ExecutionOptions, config *ExecutionConfig) e UpdatedAt: time.Now(), } - // Save execution to database + // Save execution to database (this generates ExecutionID) if err := SaveExecution(execution); err != nil { return fmt.Errorf("failed to create execution record: %w", err) } + // For ExecutionTypeFunc, register the function in global registry using ExecutionID + if config.Type == ExecutionTypeFunc && funcToRegister != nil { + config.FuncID = execution.ExecutionID // Set FuncID for later lookup + RegisterFunc(execution.ExecutionID, funcToRegister) + } + return nil } diff --git a/job/goroutine.go b/job/goroutine.go index 2062974d..67d31dd1 100644 --- a/job/goroutine.go +++ b/job/goroutine.go @@ -233,6 +233,84 @@ func UpdateExecutionProgress(executionID string, progressData map[string]interfa return nil } +// ExecuteFunc executes a Go function using goroutine mode +func (g *Goroutine) ExecuteFunc(ctx context.Context, work *WorkRequest, progress *Progress) error { + config := work.Execution.ExecutionConfig + + // Get function from global registry using FuncID (ExecutionID) + funcID := config.FuncID + if funcID == "" { + funcID = work.Execution.ExecutionID // Fallback to ExecutionID + } + + fn, ok := GetFunc(funcID) + if !ok || fn == nil { + return fmt.Errorf("execution function not found in registry (funcID: %s)", funcID) + } + + // Ensure cleanup after execution (success or failure) + defer UnregisterFunc(funcID) + + funcName := config.FuncName + if funcName == "" { + funcName = "anonymous" + } + + work.Execution.Info("Executing function: %s (goroutine mode, funcID: %s)", funcName, funcID) + + // Create execution context + execCtx := &ExecutionContext{ + Ctx: ctx, + Execution: work.Execution, + Args: config.FuncArgs, + } + + // Execute the function + err := fn(execCtx) + if err != nil { + // Check if it was cancelled + if ctx.Err() != nil { + work.Execution.Warn("Function cancelled: %s", ctx.Err().Error()) + work.Execution.Status = "cancelled" + } else { + work.Execution.Error("Function failed: %s", err.Error()) + work.Execution.Status = "failed" + + // Store error info + errorInfo := map[string]interface{}{ + "error": err.Error(), + "func_name": funcName, + } + if errorBytes, jsonErr := jsoniter.Marshal(errorInfo); jsonErr == nil { + work.Execution.ErrorInfo = (*json.RawMessage)(&errorBytes) + } + } + + // Save execution status + if saveErr := SaveExecution(work.Execution); saveErr != nil { + work.Execution.Error("Failed to save execution error: %s", saveErr.Error()) + } + + if ctx.Err() != nil { + return ctx.Err() + } + return fmt.Errorf("function execution failed: %w", err) + } + + work.Execution.Info("Function completed successfully (funcID: %s)", funcID) + + // Update execution with success result + work.Execution.Status = "completed" + work.Execution.Progress = 100 + + if saveErr := SaveExecution(work.Execution); saveErr != nil { + work.Execution.Error("Failed to save execution result: %s", saveErr.Error()) + return fmt.Errorf("failed to save execution result: %w", saveErr) + } + + return nil +} + // extractProgressData extracts progress and message from callback data func extractProgressData(data map[string]interface{}) (int, string) { var progressInt int = -1 // Default to -1 to indicate no progress value diff --git a/job/job.go b/job/job.go index fad03bdf..e16895ad 100644 --- a/job/job.go +++ b/job/job.go @@ -185,7 +185,9 @@ func DaemonAndSave(mode ModeType, data map[string]interface{}) (*Job, error) { // Push pushes the job to execution queue (renamed from Start for better semantics) func (j *Job) Push() error { - // Get executions for this job + // Get executions from database + // For ExecutionTypeFunc, the function is stored in global registry (funcRegistry) + // and will be looked up by FuncID (ExecutionID) during execution executions, err := j.GetExecutions() if err != nil { return fmt.Errorf("failed to get executions: %w", err) diff --git a/job/job_test.go b/job/job_test.go index 486f572b..8225b157 100644 --- a/job/job_test.go +++ b/job/job_test.go @@ -678,3 +678,274 @@ func TestDaemonAndSave(t *testing.T) { t.Log("DaemonAndSave job created and saved successfully") } + +// TestAddFunc tests the AddFunc method for adding Go functions as job executions +func TestAddFunc(t *testing.T) { + // Setup + test.Prepare(&testing.T{}, config.Conf) + defer test.Clean() + + // Create a job + testJob, err := job.OnceAndSave(job.GOROUTINE, map[string]interface{}{ + "name": "Test AddFunc Job", + "description": "Testing Go function execution", + }) + if err != nil { + t.Fatalf("Failed to create job: %v", err) + } + + // Track if function was called + funcCalled := false + funcArgs := make(map[string]interface{}) + + // Add a Go function execution + err = testJob.AddFunc(&job.ExecutionOptions{ + Priority: 1, + }, "test.func", func(ctx *job.ExecutionContext) error { + funcCalled = true + funcArgs = ctx.Args + t.Logf("Function executed with args: %v", ctx.Args) + return nil + }, map[string]interface{}{ + "key1": "value1", + "key2": 42, + }) + if err != nil { + t.Fatalf("Failed to add function execution: %v", err) + } + + // Get the execution to verify it was saved + executions, err := testJob.GetExecutions() + if err != nil { + t.Fatalf("Failed to get executions: %v", err) + } + if len(executions) != 1 { + t.Fatalf("Expected 1 execution, got %d", len(executions)) + } + + // Verify function is registered in global registry + funcID := executions[0].ExecutionID + fn, ok := job.GetFunc(funcID) + if !ok || fn == nil { + t.Error("Expected function to be registered in global registry") + } + + // Push the job + err = testJob.Push() + if err != nil { + t.Fatalf("Failed to push job: %v", err) + } + + // Wait for execution to complete + time.Sleep(2 * time.Second) + + // Verify function was called + if !funcCalled { + t.Error("Expected function to be called") + } + + // Verify args were passed + if funcArgs["key1"] != "value1" { + t.Errorf("Expected key1=value1, got %v", funcArgs["key1"]) + } + // Note: JSON unmarshaling converts numbers to float64 + key2Val, ok := funcArgs["key2"].(float64) + if !ok { + // Try int in case it wasn't serialized + if intVal, ok := funcArgs["key2"].(int); ok { + key2Val = float64(intVal) + } else { + t.Errorf("Expected key2 to be a number, got %T: %v", funcArgs["key2"], funcArgs["key2"]) + } + } + if key2Val != 42 { + t.Errorf("Expected key2=42, got %v", key2Val) + } + + // Verify function was cleaned up from registry after execution + fn, ok = job.GetFunc(funcID) + if ok || fn != nil { + t.Error("Expected function to be removed from global registry after execution") + } + + t.Log("AddFunc test completed successfully") +} + +// TestAddFuncMemoryCleanup tests that memory is properly cleaned up after function execution +func TestAddFuncMemoryCleanup(t *testing.T) { + // Setup + test.Prepare(&testing.T{}, config.Conf) + defer test.Clean() + + // Create a job + testJob, err := job.OnceAndSave(job.GOROUTINE, map[string]interface{}{ + "name": "Test AddFunc Memory Cleanup", + "description": "Testing memory cleanup after function execution", + }) + if err != nil { + t.Fatalf("Failed to create job: %v", err) + } + + // Create a large closure to make memory leak more detectable + largeData := make([]byte, 1024*1024) // 1MB + for i := range largeData { + largeData[i] = byte(i % 256) + } + + executed := false + + // Add a Go function with large closure + err = testJob.AddFunc(&job.ExecutionOptions{ + Priority: 1, + }, "test.cleanup", func(ctx *job.ExecutionContext) error { + // Use largeData to ensure it's captured in closure + _ = len(largeData) + executed = true + return nil + }, map[string]interface{}{ + "test": "cleanup", + }) + if err != nil { + t.Fatalf("Failed to add function execution: %v", err) + } + + // Get the execution to verify FuncID is set + executions, err := testJob.GetExecutions() + if err != nil { + t.Fatalf("Failed to get executions: %v", err) + } + if len(executions) != 1 { + t.Fatalf("Expected 1 execution, got %d", len(executions)) + } + funcID := executions[0].ExecutionID + t.Logf("FuncID (ExecutionID): %s", funcID) + + // Verify function is registered in global registry before execution + fn, ok := job.GetFunc(funcID) + if !ok || fn == nil { + t.Error("Expected function to be registered in global registry before execution") + } + + // Push the job + err = testJob.Push() + if err != nil { + t.Fatalf("Failed to push job: %v", err) + } + + // Wait for execution to complete with polling + maxWait := 10 * time.Second + pollInterval := 200 * time.Millisecond + startTime := time.Now() + + for time.Since(startTime) < maxWait { + if executed { + break + } + time.Sleep(pollInterval) + } + + // Verify function was executed + if !executed { + t.Error("Expected function to be executed") + } + + // Wait for execution to complete in database + var finalStatus string + for time.Since(startTime) < maxWait { + executions, err := testJob.GetExecutions() + if err == nil && len(executions) > 0 { + finalStatus = executions[0].Status + t.Logf("Execution status: %s", finalStatus) + if finalStatus == "completed" || finalStatus == "failed" { + break + } + } + time.Sleep(pollInterval) + } + + // Wait a bit more for cleanup to complete + time.Sleep(500 * time.Millisecond) + + // Verify memory cleanup: function should be removed from global registry + fn, ok = job.GetFunc(funcID) + if ok || fn != nil { + t.Errorf("Expected function to be removed from global registry after completion") + } + + // Verify execution status in database + executions, err = testJob.GetExecutions() + if err != nil { + t.Fatalf("Failed to get executions: %v", err) + } + + if len(executions) != 1 { + t.Errorf("Expected 1 execution in database, got %d", len(executions)) + } + + if executions[0].Status != "completed" { + t.Errorf("Expected execution status 'completed', got '%s'", executions[0].Status) + } + + t.Log("AddFunc memory cleanup test completed successfully") +} + +// TestAddFuncError tests error handling in AddFunc execution +func TestAddFuncError(t *testing.T) { + // Setup + test.Prepare(&testing.T{}, config.Conf) + defer test.Clean() + + // Create a job + testJob, err := job.OnceAndSave(job.GOROUTINE, map[string]interface{}{ + "name": "Test AddFunc Error", + "description": "Testing error handling in function execution", + }) + if err != nil { + t.Fatalf("Failed to create job: %v", err) + } + + // Add a Go function that returns an error + err = testJob.AddFunc(&job.ExecutionOptions{ + Priority: 1, + }, "test.error", func(ctx *job.ExecutionContext) error { + return fmt.Errorf("intentional test error") + }, nil) + if err != nil { + t.Fatalf("Failed to add function execution: %v", err) + } + + // Push the job + err = testJob.Push() + if err != nil { + t.Fatalf("Failed to push job: %v", err) + } + + // Wait for execution to complete + time.Sleep(2 * time.Second) + + // Verify execution failed + executions, err := testJob.GetExecutions() + if err != nil { + t.Fatalf("Failed to get executions: %v", err) + } + + if len(executions) != 1 { + t.Errorf("Expected 1 execution, got %d", len(executions)) + } + + if executions[0].Status != "failed" { + t.Errorf("Expected execution status 'failed', got '%s'", executions[0].Status) + } + + // Verify memory cleanup even on error: function should be removed from global registry + // Get the execution ID first + if len(executions) > 0 { + funcID := executions[0].ExecutionID + fn, ok := job.GetFunc(funcID) + if ok || fn != nil { + t.Errorf("Expected function to be removed from global registry after failure") + } + } + + t.Log("AddFunc error handling test completed successfully") +} diff --git a/job/types.go b/job/types.go index 05083e2d..0f0caaab 100644 --- a/job/types.go +++ b/job/types.go @@ -59,6 +59,7 @@ type ExecutionType string const ( ExecutionTypeProcess ExecutionType = "process" // Yao process (default) ExecutionTypeCommand ExecutionType = "command" // System command + ExecutionTypeFunc ExecutionType = "func" // Go function (internal use) ) // ExecutionOptions holds common execution options @@ -96,14 +97,56 @@ func (o *ExecutionOptions) AddSharedData(key string, value interface{}) *Executi return o } +// ExecutionFunc is the function signature for ExecutionTypeFunc +// The function receives the execution context and returns an error if failed +type ExecutionFunc func(ctx *ExecutionContext) error + +// ExecutionContext provides context for ExecutionFunc +type ExecutionContext struct { + Ctx context.Context // Go context + Execution *Execution // Current execution + Args map[string]interface{} // Function arguments +} + +// funcRegistry is a global registry for ExecutionFunc +// Key is the funcID (execution_id), value is the function +var funcRegistry = make(map[string]ExecutionFunc) +var funcRegistryMutex sync.RWMutex + +// RegisterFunc registers a function in the global registry +func RegisterFunc(funcID string, fn ExecutionFunc) { + funcRegistryMutex.Lock() + defer funcRegistryMutex.Unlock() + funcRegistry[funcID] = fn +} + +// GetFunc retrieves a function from the global registry +func GetFunc(funcID string) (ExecutionFunc, bool) { + funcRegistryMutex.RLock() + defer funcRegistryMutex.RUnlock() + fn, ok := funcRegistry[funcID] + return fn, ok +} + +// UnregisterFunc removes a function from the global registry +func UnregisterFunc(funcID string) { + funcRegistryMutex.Lock() + defer funcRegistryMutex.Unlock() + delete(funcRegistry, funcID) +} + // ExecutionConfig holds execution configuration based on type type ExecutionConfig struct { - Type ExecutionType `json:"type"` - ProcessName string `json:"process_name,omitempty"` // Yao process name - ProcessArgs []interface{} `json:"process_args,omitempty"` // Yao process arguments - Command string `json:"command,omitempty"` // System command - CommandArgs []string `json:"command_args,omitempty"` // Command arguments - Environment map[string]string `json:"environment,omitempty"` // Environment variables + Type ExecutionType `json:"type"` + ProcessName string `json:"process_name,omitempty"` // Yao process name + ProcessArgs []interface{} `json:"process_args,omitempty"` // Yao process arguments + Command string `json:"command,omitempty"` // System command + CommandArgs []string `json:"command_args,omitempty"` // Command arguments + Environment map[string]string `json:"environment,omitempty"` // Environment variables + Func ExecutionFunc `json:"-"` // Go function (not serialized, use FuncID instead) + FuncID string `json:"func_id,omitempty"` // Function ID for registry lookup + FuncName string `json:"func_name,omitempty"` // Function name for logging + FuncArgs map[string]interface{} `json:"func_args,omitempty"` // Function arguments } // Job represents the main job entity diff --git a/job/worker.go b/job/worker.go index 7d0b1899..6a7976a3 100644 --- a/job/worker.go +++ b/job/worker.go @@ -371,11 +371,11 @@ func (w *Worker) processWork(work *WorkRequest) { } // Clean up execution context from job + work.Job.executionMutex.Lock() if work.Job.executionContexts != nil { - work.Job.executionMutex.Lock() delete(work.Job.executionContexts, work.Execution.ExecutionID) - work.Job.executionMutex.Unlock() } + work.Job.executionMutex.Unlock() log.Debug("Worker %s finished processing job %s", w.ID, work.Job.JobID) } @@ -397,6 +397,9 @@ func (w *Worker) executeInGoroutine(ctx context.Context, work *WorkRequest, prog case ExecutionTypeCommand: return goroutineExecutor.ExecuteSystemCommand(ctx, work, progress) + case ExecutionTypeFunc: + return goroutineExecutor.ExecuteFunc(ctx, work, progress) + default: return fmt.Errorf("unsupported execution type: %s", work.Execution.ExecutionConfig.Type) } @@ -423,6 +426,11 @@ func (w *Worker) executeInProcess(ctx context.Context, work *WorkRequest, progre case ExecutionTypeCommand: return processExecutor.ExecuteSystemCommand(ctx, work, progress) + case ExecutionTypeFunc: + // ExecutionTypeFunc is not supported in process mode, fall back to goroutine + goroutineExecutor := &Goroutine{} + return goroutineExecutor.ExecuteFunc(ctx, work, progress) + default: return fmt.Errorf("unsupported execution type: %s", work.Execution.ExecutionConfig.Type) } diff --git a/kb/api/addfile.go b/kb/api/addfile.go new file mode 100644 index 00000000..8707ca92 --- /dev/null +++ b/kb/api/addfile.go @@ -0,0 +1,303 @@ +package api + +import ( + "context" + "fmt" + + "github.com/yaoapp/gou/graphrag/utils" + "github.com/yaoapp/kun/maps" + "github.com/yaoapp/yao/attachment" + "github.com/yaoapp/yao/job" +) + +// AddFile adds a file to a collection (sync) +func (instance *KBInstance) AddFile(ctx context.Context, params *AddFileParams) (*AddDocumentResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.FileID == "" { + return nil, fmt.Errorf("file_id is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Set default uploader + uploader := params.Uploader + if uploader == "" { + uploader = DefaultUploader + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Get file manager + m, ok := attachment.Managers[uploader] + if !ok { + return nil, fmt.Errorf("invalid uploader: %s not found", uploader) + } + + // Check if the file exists + exists := m.Exists(ctx, params.FileID) + if !exists { + return nil, fmt.Errorf("file not found: %s", params.FileID) + } + + // Get file info and path + path, contentType, err := m.LocalPath(ctx, params.FileID) + if err != nil { + return nil, fmt.Errorf("failed to get local path: %w", err) + } + + fileInfo, err := m.Info(ctx, params.FileID) + if err != nil { + return nil, fmt.Errorf("failed to get file info: %w", err) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": fileInfo.Filename, + "type": "file", + "status": "pending", + "uploader_id": uploader, + "file_id": params.FileID, + "file_name": fileInfo.Filename, + "file_path": path, + "file_mime_type": contentType, + "size": int64(fileInfo.Bytes), + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err = instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Process file content + params.DocID = docID // Ensure docID is set + err = instance.processFile(ctx, docID, params) + if err != nil { + return nil, err + } + + return &AddDocumentResult{ + Message: "File added successfully", + CollectionID: params.CollectionID, + DocID: docID, + FileID: params.FileID, + }, nil +} + +// AddFileAsync adds a file to a collection (async) +func (instance *KBInstance) AddFileAsync(ctx context.Context, params *AddFileParams) (*AddDocumentAsyncResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.FileID == "" { + return nil, fmt.Errorf("file_id is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Set default uploader + uploader := params.Uploader + if uploader == "" { + uploader = DefaultUploader + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Get file manager + m, ok := attachment.Managers[uploader] + if !ok { + return nil, fmt.Errorf("invalid uploader: %s not found", uploader) + } + + // Check if the file exists + exists := m.Exists(ctx, params.FileID) + if !exists { + return nil, fmt.Errorf("file not found: %s", params.FileID) + } + + // Get file info and path + path, contentType, err := m.LocalPath(ctx, params.FileID) + if err != nil { + return nil, fmt.Errorf("failed to get local path: %w", err) + } + + fileInfo, err := m.Info(ctx, params.FileID) + if err != nil { + return nil, fmt.Errorf("failed to get file info: %w", err) + } + + // Get job options with defaults + jobName, jobDescription, jobIcon, jobCategory := getJobOptions(params.Job, + "Knowledge Base File Processing", + "Processing and indexing file content for knowledge base search", + "library_add", + "Knowledge Base", + ) + + // Create job data + jobCreateData := map[string]interface{}{ + "name": jobName, + "description": jobDescription, + "category_name": jobCategory, + } + if jobIcon != "" { + jobCreateData["icon"] = jobIcon + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + jobCreateData[k] = v + } + } + + // Create and save Job + j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + if err != nil { + return nil, fmt.Errorf("failed to create and save job: %w", err) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": fileInfo.Filename, + "type": "file", + "status": "pending", + "uploader_id": uploader, + "file_id": params.FileID, + "file_name": fileInfo.Filename, + "file_path": path, + "file_mime_type": contentType, + "size": int64(fileInfo.Bytes), + "job_id": j.JobID, + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err = instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Capture parameters for the async function + asyncDocID := docID + asyncParams := &AddFileParams{ + CollectionID: params.CollectionID, + FileID: params.FileID, + Uploader: uploader, + Locale: params.Locale, + Chunking: params.Chunking, + Embedding: params.Embedding, + Extraction: params.Extraction, + Fetcher: params.Fetcher, + Converter: params.Converter, + } + + // Add function execution to job + err = j.AddFunc(&job.ExecutionOptions{Priority: 1}, "kb.addfile", func(execCtx *job.ExecutionContext) error { + return instance.processFile(execCtx.Ctx, asyncDocID, asyncParams) + }, map[string]interface{}{ + "doc_id": asyncDocID, + "collection_id": params.CollectionID, + "file_id": params.FileID, + }) + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to add job execution: %w", err) + } + + // Push the job to execution queue + err = j.Push() + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to push job: %w", err) + } + + return &AddDocumentAsyncResult{ + JobID: j.JobID, + DocID: docID, + }, nil +} + +// processFile processes file content and updates the knowledge base +func (instance *KBInstance) processFile(ctx context.Context, docID string, params *AddFileParams) error { + uploader := params.Uploader + if uploader == "" { + uploader = DefaultUploader + } + + // Get file manager + m, ok := attachment.Managers[uploader] + if !ok { + return fmt.Errorf("invalid uploader: %s not found", uploader) + } + + // Get file path + path, contentType, err := m.LocalPath(ctx, params.FileID) + if err != nil { + return fmt.Errorf("failed to get local path: %w", err) + } + + // Convert to UpsertOptions + upsertOptions, err := instance.toUpsertOptions(docID, params.CollectionID, params.Locale, path, contentType, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to convert to upsert options: %w", err) + } + + // Add file to GraphRag + _, err = instance.GraphRag.AddFile(ctx, path, upsertOptions) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to add file: %w", err) + } + + // Update status and segment count + instance.updateDocumentAfterProcessing(ctx, docID, params.CollectionID) + + return nil +} diff --git a/kb/api/addfile_test.go b/kb/api/addfile_test.go new file mode 100644 index 00000000..1178427a --- /dev/null +++ b/kb/api/addfile_test.go @@ -0,0 +1,585 @@ +package api_test + +import ( + "context" + "fmt" + "mime/multipart" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/attachment" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go, which handles environment setup +// Run tests with: source env.local.sh && go test -v ./kb/api/... + +// createTestCollectionForFile is a helper to create a test collection for file tests +func createTestCollectionForFile(t *testing.T, ctx context.Context) string { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + collectionID := fmt.Sprintf("test_file_%d", time.Now().UnixNano()) + + params := &api.CreateCollectionParams{ + ID: collectionID, + Metadata: map[string]interface{}{ + "name": "Test File Collection", + "description": "Collection for AddFile tests", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + + _, err := kb.API.CreateCollection(ctx, params) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + + return collectionID +} + +// cleanupTestCollectionForFile removes a test collection +func cleanupTestCollectionForFile(ctx context.Context, collectionID string) { + if kb.API != nil { + _, _ = kb.API.RemoveCollection(ctx, collectionID) + } +} + +// ========== AddFile Tests ========== +// Note: Full AddFile tests require actual files to be uploaded via attachment manager +// These tests verify parameter validation and error handling + +func TestAddFile(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForFile(t, ctx) + defer cleanupTestCollectionForFile(ctx, collectionID) + + t.Run("AddFileMissingCollectionID", func(t *testing.T) { + params := &api.AddFileParams{ + FileID: "some_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddFileMissingFileID", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "file_id is required") + }) + + t.Run("AddFileMissingChunking", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddFileMissingEmbedding", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddFileInvalidUploader", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Uploader: "invalid_uploader", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "invalid uploader") + }) + + t.Run("AddFileNotFound", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "nonexistent_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + // Error could be "file not found" or "invalid uploader" depending on environment + assert.True(t, err != nil, "Expected an error") + }) +} + +// ========== AddFileAsync Tests ========== + +func TestAddFileAsync(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForFile(t, ctx) + defer cleanupTestCollectionForFile(ctx, collectionID) + + t.Run("AddFileAsyncMissingCollectionID", func(t *testing.T) { + params := &api.AddFileParams{ + FileID: "some_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddFileAsyncMissingFileID", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "file_id is required") + }) + + t.Run("AddFileAsyncMissingChunking", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddFileAsyncMissingEmbedding", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddFileAsyncInvalidUploader", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "some_file_id", + Uploader: "invalid_uploader", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "invalid uploader") + }) + + t.Run("AddFileAsyncNotFound", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: "nonexistent_file_id", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + // Error could be "file not found" or "invalid uploader" depending on environment + assert.True(t, err != nil, "Expected an error") + }) +} + +// ========== AddFile with Real File Tests ========== + +// getTestUploader returns the uploader name and manager for testing +func getTestUploader(t *testing.T) (string, *attachment.Manager) { + // Try __yao.attachment first (system uploader) + if manager, ok := attachment.Managers["__yao.attachment"]; ok { + return "__yao.attachment", manager + } + // Try local manager + if manager, ok := attachment.Managers["local"]; ok { + return "local", manager + } + // List available managers for debugging + var available []string + for name := range attachment.Managers { + available = append(available, name) + } + t.Fatalf("No attachment manager available. Available managers: %v", available) + return "", nil +} + +// uploadTestFile uploads a test file using the attachment manager and returns the file ID +func uploadTestFile(t *testing.T, ctx context.Context, filename, content string) string { + _, manager := getTestUploader(t) + + // Create file header + fileHeader := &attachment.FileHeader{ + FileHeader: &multipart.FileHeader{ + Filename: filename, + Size: int64(len(content)), + Header: make(map[string][]string), + }, + } + fileHeader.Header.Set("Content-Type", "text/plain") + + // Upload the file + reader := strings.NewReader(content) + file, err := manager.Upload(ctx, fileHeader, reader, attachment.UploadOption{}) + if err != nil { + t.Fatalf("Failed to upload test file: %v", err) + } + + t.Logf("Uploaded test file: %s (ID: %s)", filename, file.ID) + return file.ID +} + +// cleanupTestFile removes a test file +func cleanupTestFile(ctx context.Context, t *testing.T, fileID string) { + _, manager := getTestUploader(t) + _ = manager.Delete(ctx, fileID) +} + +func TestAddFileWithRealFile(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForFile(t, ctx) + defer cleanupTestCollectionForFile(ctx, collectionID) + + // Get the uploader name + uploaderName, _ := getTestUploader(t) + + // Upload a test file + testContent := `This is a test document for the knowledge base. +It contains content to test the file processing functionality.` + + fileID := uploadTestFile(t, ctx, "test_document.txt", testContent) + defer cleanupTestFile(ctx, t, fileID) + + t.Run("AddFileSuccess", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: fileID, + Uploader: uploaderName, + Locale: "en", + Metadata: map[string]interface{}{ + "description": "A test file document", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "test_user", + }, + } + + result, err := kb.API.AddFile(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, collectionID, result.CollectionID) + assert.NotEmpty(t, result.DocID) + assert.Equal(t, fileID, result.FileID) + assert.Contains(t, result.Message, "successfully") + t.Logf("Added file document: %s", result.DocID) + } + + // Verify document was created + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "file", doc["type"]) + assert.Equal(t, "completed", doc["status"]) + t.Logf("✅ File Document verified: type=%v, status=%v", doc["type"], doc["status"]) + } + }) +} + +func TestAddFileAsyncWithRealFile(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForFile(t, ctx) + defer cleanupTestCollectionForFile(ctx, collectionID) + + // Get the uploader name + uploaderName, _ := getTestUploader(t) + + // Upload a test file for async processing + testContent := `Async test document content. + +This document will be processed asynchronously. + +The job system should handle the processing in the background.` + + fileID := uploadTestFile(t, ctx, "async_test_document.txt", testContent) + defer cleanupTestFile(ctx, t, fileID) + + t.Run("AddFileAsyncSuccess", func(t *testing.T) { + params := &api.AddFileParams{ + CollectionID: collectionID, + FileID: fileID, + Uploader: uploaderName, + Locale: "en", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Job: &api.JobOptionsParams{ + Name: "Test Async File Job", + Description: "Testing async file processing", + Category: "Test", + }, + } + + result, err := kb.API.AddFileAsync(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.NotEmpty(t, result.JobID) + assert.NotEmpty(t, result.DocID) + t.Logf("Created async file job: %s for document: %s", result.JobID, result.DocID) + } + + // Verify document was created with pending status + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "file", doc["type"]) + assert.Equal(t, result.JobID, doc["job_id"]) + t.Logf("✅ Async file document created: status=%v, job_id=%v", doc["status"], doc["job_id"]) + + // Wait for job to complete (max 30 seconds) + maxWait := 30 * time.Second + pollInterval := 500 * time.Millisecond + startTime := time.Now() + var finalStatus string + + for time.Since(startTime) < maxWait { + doc, err = kb.API.GetDocument(ctx, result.DocID, nil) + if err != nil { + t.Logf("Error getting document: %v", err) + break + } + finalStatus, _ = doc["status"].(string) + if finalStatus == "completed" || finalStatus == "error" { + break + } + time.Sleep(pollInterval) + } + + t.Logf("✅ Job completed: final status=%s, elapsed=%v", finalStatus, time.Since(startTime)) + assert.Equal(t, "completed", finalStatus, "Job should complete successfully") + } + }) +} + +func TestAddFileIntegration(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForFile(t, ctx) + defer cleanupTestCollectionForFile(ctx, collectionID) + + // Get the uploader name + uploaderName, _ := getTestUploader(t) + + t.Run("FullFileLifecycle", func(t *testing.T) { + // Upload a test file + testContent := `Integration test document. + +This document tests the full lifecycle of file processing: +1. Upload file +2. Add to knowledge base +3. Verify document creation +4. List documents +5. Remove document + +End of test content.` + + fileID := uploadTestFile(t, ctx, "lifecycle_test.txt", testContent) + defer cleanupTestFile(ctx, t, fileID) + + // 1. Add File Document + addParams := &api.AddFileParams{ + CollectionID: collectionID, + FileID: fileID, + Uploader: uploaderName, + Locale: "en", + Metadata: map[string]interface{}{ + "title": "File Lifecycle Test", + "description": "Full lifecycle integration test", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "integration_test", + }, + } + + result, err := kb.API.AddFile(ctx, addParams) + assert.NoError(t, err) + assert.NotNil(t, result) + if result == nil { + t.Fatalf("Failed to create file document: result is nil") + } + t.Logf("1. Created file document: %s", result.DocID) + + // 2. Get Document + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "file", doc["type"]) + assert.Equal(t, "completed", doc["status"]) + t.Logf("2. Retrieved document: name=%v, type=%v, status=%v", doc["name"], doc["type"], doc["status"]) + + // 3. List Documents + listFilter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + } + listResult, err := kb.API.ListDocuments(ctx, listFilter) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(listResult.Data), 1) + t.Logf("3. Found document in list: %d documents", len(listResult.Data)) + + // 4. Remove Document + removeParams := &api.RemoveDocumentsParams{ + DocumentIDs: []string{result.DocID}, + } + removeResult, err := kb.API.RemoveDocuments(ctx, removeParams) + assert.NoError(t, err) + assert.NotNil(t, removeResult) + t.Logf("4. Removed document: %d deleted", removeResult.DeletedCount) + + // 5. Verify Removal + _, err = kb.API.GetDocument(ctx, result.DocID, nil) + assert.Error(t, err) + t.Logf("5. Verified document removal") + + t.Logf("✅ Full file lifecycle test completed successfully") + }) +} diff --git a/kb/api/addtext.go b/kb/api/addtext.go new file mode 100644 index 00000000..dbfd71e4 --- /dev/null +++ b/kb/api/addtext.go @@ -0,0 +1,230 @@ +package api + +import ( + "context" + "fmt" + + "github.com/yaoapp/gou/graphrag/utils" + "github.com/yaoapp/kun/maps" + "github.com/yaoapp/yao/job" +) + +// AddText adds text to a collection (sync) +func (instance *KBInstance) AddText(ctx context.Context, params *AddTextParams) (*AddDocumentResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.Text == "" { + return nil, fmt.Errorf("text is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": "Text Document", + "type": "text", + "status": "pending", + "text_content": params.Text, + "size": int64(len(params.Text)), + } + + // Use title from metadata if available + if params.Metadata != nil { + if title, ok := params.Metadata["title"].(string); ok && title != "" { + documentData["name"] = title + } + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err := instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Process text content + params.DocID = docID // Ensure docID is set + err = instance.processText(ctx, docID, params) + if err != nil { + return nil, err + } + + return &AddDocumentResult{ + Message: "Text added successfully", + CollectionID: params.CollectionID, + DocID: docID, + }, nil +} + +// AddTextAsync adds text to a collection (async) +func (instance *KBInstance) AddTextAsync(ctx context.Context, params *AddTextParams) (*AddDocumentAsyncResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.Text == "" { + return nil, fmt.Errorf("text is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Get job options with defaults + jobName, jobDescription, jobIcon, jobCategory := getJobOptions(params.Job, + "Knowledge Base Text Processing", + "Processing and indexing text content for knowledge base search", + "library_add", + "Knowledge Base", + ) + + // Create job data + jobCreateData := map[string]interface{}{ + "name": jobName, + "description": jobDescription, + "category_name": jobCategory, + } + if jobIcon != "" { + jobCreateData["icon"] = jobIcon + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + jobCreateData[k] = v + } + } + + // Create and save Job + j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + if err != nil { + return nil, fmt.Errorf("failed to create and save job: %w", err) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": "Text Document", + "type": "text", + "status": "pending", + "text_content": params.Text, + "size": int64(len(params.Text)), + "job_id": j.JobID, + } + + // Use title from metadata if available + if params.Metadata != nil { + if title, ok := params.Metadata["title"].(string); ok && title != "" { + documentData["name"] = title + } + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err = instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Capture parameters for the async function + asyncDocID := docID + asyncParams := &AddTextParams{ + CollectionID: params.CollectionID, + Text: params.Text, + Locale: params.Locale, + Chunking: params.Chunking, + Embedding: params.Embedding, + Extraction: params.Extraction, + Fetcher: params.Fetcher, + Converter: params.Converter, + } + + // Add function execution to job + err = j.AddFunc(&job.ExecutionOptions{Priority: 1}, "kb.addtext", func(execCtx *job.ExecutionContext) error { + return instance.processText(execCtx.Ctx, asyncDocID, asyncParams) + }, map[string]interface{}{ + "doc_id": asyncDocID, + "collection_id": params.CollectionID, + }) + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to add job execution: %w", err) + } + + // Push the job to execution queue + err = j.Push() + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to push job: %w", err) + } + + return &AddDocumentAsyncResult{ + JobID: j.JobID, + DocID: docID, + }, nil +} + +// processText processes text content and updates the knowledge base +func (instance *KBInstance) processText(ctx context.Context, docID string, params *AddTextParams) error { + // Convert to UpsertOptions + upsertOptions, err := instance.toUpsertOptions(docID, params.CollectionID, params.Locale, "", "", params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to convert to upsert options: %w", err) + } + + // Add text to GraphRag + _, err = instance.GraphRag.AddText(ctx, params.Text, upsertOptions) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to add text: %w", err) + } + + // Update status and segment count + instance.updateDocumentAfterProcessing(ctx, docID, params.CollectionID) + + return nil +} diff --git a/kb/api/addtext_test.go b/kb/api/addtext_test.go new file mode 100644 index 00000000..c8fb4ef3 --- /dev/null +++ b/kb/api/addtext_test.go @@ -0,0 +1,485 @@ +package api_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go, which handles environment setup +// Run tests with: source env.local.sh && go test -v ./kb/api/... + +// createTestCollectionForText is a helper to create a test collection for text tests +func createTestCollectionForText(t *testing.T, ctx context.Context) string { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + collectionID := fmt.Sprintf("test_text_%d", time.Now().UnixNano()) + + params := &api.CreateCollectionParams{ + ID: collectionID, + Metadata: map[string]interface{}{ + "name": "Test Text Collection", + "description": "Collection for AddText tests", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + + _, err := kb.API.CreateCollection(ctx, params) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + + return collectionID +} + +// cleanupTestCollectionForText removes a test collection +func cleanupTestCollectionForText(ctx context.Context, collectionID string) { + if kb.API != nil { + _, _ = kb.API.RemoveCollection(ctx, collectionID) + } +} + +// ========== AddText Tests ========== + +func TestAddText(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForText(t, ctx) + defer cleanupTestCollectionForText(ctx, collectionID) + + t.Run("AddTextSuccess", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "This is a test document content for knowledge base testing. It contains some sample text that will be chunked and embedded.", + Locale: "en", + Metadata: map[string]interface{}{ + "title": "Test Text Document", + "description": "A test document", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "test_user", + }, + } + + result, err := kb.API.AddText(ctx, params) + if err != nil { + // Skip if connector not loaded (environment issue) + if assert.Contains(t, err.Error(), "connector") { + t.Skipf("Skipping due to connector not loaded: %v", err) + } + } + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, collectionID, result.CollectionID) + assert.NotEmpty(t, result.DocID) + assert.Contains(t, result.Message, "successfully") + t.Logf("Added text document: %s", result.DocID) + + // Verify document was created + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "text", doc["type"]) + assert.Equal(t, "Test Text Document", doc["name"]) + assert.Equal(t, "completed", doc["status"]) + t.Logf("✅ Document verified: type=%v, name=%v, status=%v", doc["type"], doc["name"], doc["status"]) + } + }) + + t.Run("AddTextMissingCollectionID", func(t *testing.T) { + params := &api.AddTextParams{ + Text: "Some text content", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddText(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddTextMissingText", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddText(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "text is required") + }) + + t.Run("AddTextMissingChunking", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "Some text content", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddText(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddTextMissingEmbedding", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "Some text content", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddText(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddTextWithCustomDocID", func(t *testing.T) { + customDocID := fmt.Sprintf("custom_text_doc_%d", time.Now().UnixNano()) + params := &api.AddTextParams{ + CollectionID: collectionID, + DocID: customDocID, + Text: "Text with custom document ID", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + } + + result, err := kb.API.AddText(ctx, params) + if err != nil { + t.Skipf("Skipping due to error: %v", err) + } + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, customDocID, result.DocID) + t.Logf("Added text with custom DocID: %s", result.DocID) + } + }) + + t.Run("AddTextWithTitleFromMetadata", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "Text content with title from metadata", + Metadata: map[string]interface{}{ + "title": "Custom Title From Metadata", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + } + + result, err := kb.API.AddText(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.Equal(t, "Custom Title From Metadata", doc["name"]) + t.Logf("✅ Title from metadata verified: %v", doc["name"]) + } + }) +} + +// ========== AddTextAsync Tests ========== + +func TestAddTextAsync(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForText(t, ctx) + defer cleanupTestCollectionForText(ctx, collectionID) + + t.Run("AddTextAsyncSuccess", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "This is async text content for testing background processing.", + Locale: "en", + Metadata: map[string]interface{}{ + "title": "Async Text Document", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Job: &api.JobOptionsParams{ + Name: "Test Async Text Job", + Description: "Testing async text processing", + Category: "Test", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.NotEmpty(t, result.JobID) + assert.NotEmpty(t, result.DocID) + t.Logf("Created async job: %s for document: %s", result.JobID, result.DocID) + } + + // Verify document was created with pending status + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "text", doc["type"]) + assert.Equal(t, result.JobID, doc["job_id"]) + t.Logf("✅ Async document created: status=%v, job_id=%v", doc["status"], doc["job_id"]) + + // Wait for job to complete (max 30 seconds) + maxWait := 30 * time.Second + pollInterval := 500 * time.Millisecond + startTime := time.Now() + var finalStatus string + + for time.Since(startTime) < maxWait { + doc, err = kb.API.GetDocument(ctx, result.DocID, nil) + if err != nil { + t.Logf("Error getting document: %v", err) + break + } + finalStatus, _ = doc["status"].(string) + if finalStatus == "completed" || finalStatus == "error" { + break + } + time.Sleep(pollInterval) + } + + t.Logf("✅ Job completed: final status=%s, elapsed=%v", finalStatus, time.Since(startTime)) + if finalStatus == "error" { + if errMsg, ok := doc["error_message"].(string); ok { + t.Logf("Error message: %s", errMsg) + } + } + assert.Equal(t, "completed", finalStatus, "Job should complete successfully") + } + }) + + t.Run("AddTextAsyncMissingCollectionID", func(t *testing.T) { + params := &api.AddTextParams{ + Text: "Some text", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddTextAsyncMissingText", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "text is required") + }) + + t.Run("AddTextAsyncMissingChunking", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "Some text", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddTextAsyncMissingEmbedding", func(t *testing.T) { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: "Some text", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddTextAsyncWithCustomDocID", func(t *testing.T) { + customDocID := fmt.Sprintf("async_custom_text_%d", time.Now().UnixNano()) + params := &api.AddTextParams{ + CollectionID: collectionID, + DocID: customDocID, + Text: "Async text with custom DocID", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + } + + result, err := kb.API.AddTextAsync(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, customDocID, result.DocID) + t.Logf("Created async text with custom DocID: %s", result.DocID) + } + }) +} + +// ========== AddText Integration Test ========== + +func TestAddTextIntegration(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForText(t, ctx) + defer cleanupTestCollectionForText(ctx, collectionID) + + t.Run("FullTextLifecycle", func(t *testing.T) { + // 1. Add Text Document + addParams := &api.AddTextParams{ + CollectionID: collectionID, + Text: "This is a comprehensive test of the text document lifecycle including creation, retrieval, and removal.", + Locale: "en", + Metadata: map[string]interface{}{ + "title": "Text Lifecycle Test", + "description": "Full lifecycle integration test", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "integration_test", + }, + } + + result, err := kb.API.AddText(ctx, addParams) + if err != nil { + t.Skipf("Skipping integration test due to AddText error: %v", err) + return + } + assert.NotNil(t, result) + t.Logf("1. Created text document: %s", result.DocID) + + // 2. Get Document + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "Text Lifecycle Test", doc["name"]) + assert.Equal(t, "text", doc["type"]) + assert.Equal(t, "completed", doc["status"]) + t.Logf("2. Retrieved document: name=%v, status=%v", doc["name"], doc["status"]) + + // 3. List Documents + listFilter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + Keywords: "Text Lifecycle", + } + listResult, err := kb.API.ListDocuments(ctx, listFilter) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(listResult.Data), 1) + t.Logf("3. Found document in list: %d documents", len(listResult.Data)) + + // 4. Remove Document + removeParams := &api.RemoveDocumentsParams{ + DocumentIDs: []string{result.DocID}, + } + removeResult, err := kb.API.RemoveDocuments(ctx, removeParams) + assert.NoError(t, err) + assert.NotNil(t, removeResult) + t.Logf("4. Removed document: %d deleted", removeResult.DeletedCount) + + // 5. Verify Removal + _, err = kb.API.GetDocument(ctx, result.DocID, nil) + assert.Error(t, err) + t.Logf("5. Verified document removal") + + t.Logf("✅ Full text lifecycle test completed successfully") + }) +} diff --git a/kb/api/addurl.go b/kb/api/addurl.go new file mode 100644 index 00000000..6eddcad7 --- /dev/null +++ b/kb/api/addurl.go @@ -0,0 +1,230 @@ +package api + +import ( + "context" + "fmt" + + "github.com/yaoapp/gou/graphrag/utils" + "github.com/yaoapp/kun/maps" + "github.com/yaoapp/yao/job" +) + +// AddURL adds a URL to a collection (sync) +func (instance *KBInstance) AddURL(ctx context.Context, params *AddURLParams) (*AddDocumentResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.URL == "" { + return nil, fmt.Errorf("url is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": "URL Document", + "type": "url", + "status": "pending", + "url": params.URL, + } + + // Use title from metadata if available + if params.Metadata != nil { + if title, ok := params.Metadata["title"].(string); ok && title != "" { + documentData["name"] = title + } + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err := instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Process URL content + params.DocID = docID // Ensure docID is set + err = instance.processURL(ctx, docID, params) + if err != nil { + return nil, err + } + + return &AddDocumentResult{ + Message: "URL added successfully", + CollectionID: params.CollectionID, + DocID: docID, + URL: params.URL, + }, nil +} + +// AddURLAsync adds a URL to a collection (async) +func (instance *KBInstance) AddURLAsync(ctx context.Context, params *AddURLParams) (*AddDocumentAsyncResult, error) { + // Validate required parameters + if params.CollectionID == "" { + return nil, fmt.Errorf("collection_id is required") + } + if params.URL == "" { + return nil, fmt.Errorf("url is required") + } + if params.Chunking == nil { + return nil, fmt.Errorf("chunking configuration is required") + } + if params.Embedding == nil { + return nil, fmt.Errorf("embedding configuration is required") + } + + // Generate document ID if not provided + docID := params.DocID + if docID == "" { + docID = utils.GenDocIDWithCollectionID(params.CollectionID) + } + + // Get job options with defaults + jobName, jobDescription, jobIcon, jobCategory := getJobOptions(params.Job, + "Knowledge Base Web Content Processing", + "Fetching and indexing web content for knowledge base search", + "library_add", + "Knowledge Base", + ) + + // Create job data + jobCreateData := map[string]interface{}{ + "name": jobName, + "description": jobDescription, + "category_name": jobCategory, + } + if jobIcon != "" { + jobCreateData["icon"] = jobIcon + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + jobCreateData[k] = v + } + } + + // Create and save Job + j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + if err != nil { + return nil, fmt.Errorf("failed to create and save job: %w", err) + } + + // Create document record + documentData := map[string]interface{}{ + "document_id": docID, + "collection_id": params.CollectionID, + "name": "URL Document", + "type": "url", + "status": "pending", + "url": params.URL, + "job_id": j.JobID, + } + + // Use title from metadata if available + if params.Metadata != nil { + if title, ok := params.Metadata["title"].(string); ok && title != "" { + documentData["name"] = title + } + } + + // Add auth scope fields + if params.AuthScope != nil { + for k, v := range params.AuthScope { + documentData[k] = v + } + } + + // Add base fields + addBaseFieldsFromParams(documentData, params.Locale, params.Metadata, params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + + // Create database record + _, err = instance.Config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + return nil, fmt.Errorf("failed to save document metadata: %w", err) + } + + // Capture parameters for the async function + asyncDocID := docID + asyncParams := &AddURLParams{ + CollectionID: params.CollectionID, + URL: params.URL, + Locale: params.Locale, + Chunking: params.Chunking, + Embedding: params.Embedding, + Extraction: params.Extraction, + Fetcher: params.Fetcher, + Converter: params.Converter, + } + + // Add function execution to job + err = j.AddFunc(&job.ExecutionOptions{Priority: 1}, "kb.addurl", func(execCtx *job.ExecutionContext) error { + return instance.processURL(execCtx.Ctx, asyncDocID, asyncParams) + }, map[string]interface{}{ + "doc_id": asyncDocID, + "collection_id": params.CollectionID, + "url": params.URL, + }) + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to add job execution: %w", err) + } + + // Push the job to execution queue + err = j.Push() + if err != nil { + // Rollback: remove document record + instance.Config.RemoveDocument(docID) + return nil, fmt.Errorf("failed to push job: %w", err) + } + + return &AddDocumentAsyncResult{ + JobID: j.JobID, + DocID: docID, + }, nil +} + +// processURL processes URL content and updates the knowledge base +func (instance *KBInstance) processURL(ctx context.Context, docID string, params *AddURLParams) error { + // Convert to UpsertOptions + upsertOptions, err := instance.toUpsertOptions(docID, params.CollectionID, params.Locale, "", "", params.Chunking, params.Embedding, params.Extraction, params.Fetcher, params.Converter) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to convert to upsert options: %w", err) + } + + // Add URL to GraphRag + _, err = instance.GraphRag.AddURL(ctx, params.URL, upsertOptions) + if err != nil { + instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + return fmt.Errorf("failed to add URL: %w", err) + } + + // Update status and segment count + instance.updateDocumentAfterProcessing(ctx, docID, params.CollectionID) + + return nil +} diff --git a/kb/api/addurl_test.go b/kb/api/addurl_test.go new file mode 100644 index 00000000..fc48827f --- /dev/null +++ b/kb/api/addurl_test.go @@ -0,0 +1,499 @@ +package api_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go, which handles environment setup +// Run tests with: source env.local.sh && go test -v ./kb/api/... + +// createTestCollectionForURL is a helper to create a test collection for URL tests +func createTestCollectionForURL(t *testing.T, ctx context.Context) string { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + collectionID := fmt.Sprintf("test_url_%d", time.Now().UnixNano()) + + params := &api.CreateCollectionParams{ + ID: collectionID, + Metadata: map[string]interface{}{ + "name": "Test URL Collection", + "description": "Collection for AddURL tests", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + + _, err := kb.API.CreateCollection(ctx, params) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + + return collectionID +} + +// cleanupTestCollectionForURL removes a test collection +func cleanupTestCollectionForURL(ctx context.Context, collectionID string) { + if kb.API != nil { + _, _ = kb.API.RemoveCollection(ctx, collectionID) + } +} + +// ========== AddURL Tests ========== + +func TestAddURL(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForURL(t, ctx) + defer cleanupTestCollectionForURL(ctx, collectionID) + + t.Run("AddURLSuccess", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Locale: "en", + Metadata: map[string]interface{}{ + "title": "Example Website", + "description": "A test URL document", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "test_user", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, collectionID, result.CollectionID) + assert.NotEmpty(t, result.DocID) + assert.Equal(t, "https://example.com", result.URL) + assert.Contains(t, result.Message, "successfully") + t.Logf("Added URL document: %s", result.DocID) + } + + // Verify document was created + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "url", doc["type"]) + assert.Equal(t, "https://example.com", doc["url"]) + t.Logf("✅ URL Document verified: type=%v, url=%v, status=%v", doc["type"], doc["url"], doc["status"]) + } + }) + + t.Run("AddURLMissingCollectionID", func(t *testing.T) { + params := &api.AddURLParams{ + URL: "https://example.com", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddURLMissingURL", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "url is required") + }) + + t.Run("AddURLMissingChunking", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddURLMissingEmbedding", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddURLWithCustomDocID", func(t *testing.T) { + customDocID := fmt.Sprintf("custom_url_doc_%d", time.Now().UnixNano()) + params := &api.AddURLParams{ + CollectionID: collectionID, + DocID: customDocID, + URL: "https://example.com", // Use root URL which always exists + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, customDocID, result.DocID) + t.Logf("Added URL with custom DocID: %s", result.DocID) + } + }) + + t.Run("AddURLWithTitleFromMetadata", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", // Use root URL which always exists + Metadata: map[string]interface{}{ + "title": "Custom URL Title From Metadata", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + } + + result, err := kb.API.AddURL(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.Equal(t, "Custom URL Title From Metadata", doc["name"]) + t.Logf("✅ Title from metadata verified: %v", doc["name"]) + } + }) +} + +// ========== AddURLAsync Tests ========== + +func TestAddURLAsync(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForURL(t, ctx) + defer cleanupTestCollectionForURL(ctx, collectionID) + + t.Run("AddURLAsyncSuccess", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Locale: "en", + Metadata: map[string]interface{}{ + "title": "Async URL Document", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + Job: &api.JobOptionsParams{ + Name: "Test Async URL Job", + Description: "Testing async URL processing", + Category: "Test", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.NotEmpty(t, result.JobID) + assert.NotEmpty(t, result.DocID) + t.Logf("Created async URL job: %s for document: %s", result.JobID, result.DocID) + } + + // Verify document was created + if result != nil { + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "url", doc["type"]) + assert.Equal(t, result.JobID, doc["job_id"]) + t.Logf("✅ Async URL document created: status=%v, job_id=%v", doc["status"], doc["job_id"]) + + // Wait for job to complete (max 30 seconds) + maxWait := 30 * time.Second + pollInterval := 500 * time.Millisecond + startTime := time.Now() + var finalStatus string + + for time.Since(startTime) < maxWait { + doc, err = kb.API.GetDocument(ctx, result.DocID, nil) + if err != nil { + t.Logf("Error getting document: %v", err) + break + } + finalStatus, _ = doc["status"].(string) + if finalStatus == "completed" || finalStatus == "error" { + break + } + time.Sleep(pollInterval) + } + + t.Logf("✅ Job completed: final status=%s, elapsed=%v", finalStatus, time.Since(startTime)) + assert.Equal(t, "completed", finalStatus, "Job should complete successfully") + } + }) + + t.Run("AddURLAsyncMissingCollectionID", func(t *testing.T) { + params := &api.AddURLParams{ + URL: "https://example.com", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "collection_id is required") + }) + + t.Run("AddURLAsyncMissingURL", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "url is required") + }) + + t.Run("AddURLAsyncMissingChunking", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "chunking configuration is required") + }) + + t.Run("AddURLAsyncMissingEmbedding", func(t *testing.T) { + params := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "embedding configuration is required") + }) + + t.Run("AddURLAsyncWithCustomDocID", func(t *testing.T) { + customDocID := fmt.Sprintf("async_custom_url_%d", time.Now().UnixNano()) + params := &api.AddURLParams{ + CollectionID: collectionID, + DocID: customDocID, + URL: "https://example.com/async-custom", + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + } + + result, err := kb.API.AddURLAsync(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + if result != nil { + assert.Equal(t, customDocID, result.DocID) + t.Logf("Created async URL with custom DocID: %s", result.DocID) + } + }) +} + +// ========== AddURL Integration Test ========== + +func TestAddURLIntegration(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForURL(t, ctx) + defer cleanupTestCollectionForURL(ctx, collectionID) + + t.Run("FullURLLifecycle", func(t *testing.T) { + // 1. Add URL Document + addParams := &api.AddURLParams{ + CollectionID: collectionID, + URL: "https://example.com", // Use root URL which always exists + Locale: "en", + Metadata: map[string]interface{}{ + "title": "URL Lifecycle Test", + "description": "Full lifecycle integration test", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Fetcher: &api.ProviderConfigParams{ + ProviderID: "__yao.http", + OptionID: "http", + }, + AuthScope: map[string]interface{}{ + "__yao_created_by": "integration_test", + }, + } + + result, err := kb.API.AddURL(ctx, addParams) + assert.NoError(t, err) + assert.NotNil(t, result) + if result == nil { + t.Fatalf("Failed to create URL document: result is nil") + } + t.Logf("1. Created URL document: %s", result.DocID) + + // 2. Get Document + doc, err := kb.API.GetDocument(ctx, result.DocID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, "URL Lifecycle Test", doc["name"]) + assert.Equal(t, "url", doc["type"]) + assert.Equal(t, "https://example.com", doc["url"]) + t.Logf("2. Retrieved document: name=%v, url=%v, status=%v", doc["name"], doc["url"], doc["status"]) + + // 3. List Documents + listFilter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + Keywords: "URL Lifecycle", + } + listResult, err := kb.API.ListDocuments(ctx, listFilter) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(listResult.Data), 1) + t.Logf("3. Found document in list: %d documents", len(listResult.Data)) + + // 4. Remove Document + removeParams := &api.RemoveDocumentsParams{ + DocumentIDs: []string{result.DocID}, + } + removeResult, err := kb.API.RemoveDocuments(ctx, removeParams) + assert.NoError(t, err) + assert.NotNil(t, removeResult) + t.Logf("4. Removed document: %d deleted", removeResult.DeletedCount) + + // 5. Verify Removal + _, err = kb.API.GetDocument(ctx, result.DocID, nil) + assert.Error(t, err) + t.Logf("5. Verified document removal") + + t.Logf("✅ Full URL lifecycle test completed successfully") + }) +} diff --git a/kb/api/collection_test.go b/kb/api/collection_test.go index 17332cee..ed440771 100644 --- a/kb/api/collection_test.go +++ b/kb/api/collection_test.go @@ -10,6 +10,7 @@ import ( "github.com/stretchr/testify/assert" graphragtypes "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/gou/model" + "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/kb/api" @@ -21,8 +22,14 @@ func TestMain(m *testing.M) { test.Prepare(&testing.T{}, config.Conf) defer test.Clean() + // Load attachment managers (needed for file upload tests) + err := attachment.Load(config.Conf) + if err != nil { + panic("Failed to load attachment managers: " + err.Error()) + } + // Load knowledge base - _, err := kb.Load(config.Conf) + _, err = kb.Load(config.Conf) if err != nil { panic("Failed to load knowledge base: " + err.Error()) } diff --git a/kb/api/consts.go b/kb/api/consts.go index 1520aa94..11c6ec30 100644 --- a/kb/api/consts.go +++ b/kb/api/consts.go @@ -55,3 +55,51 @@ var DefaultSort = []model.QueryOrder{ const ( DefaultLocale = "en" ) + +// Document field definitions +var ( + // AvailableDocumentFields defines all available fields for security filtering + AvailableDocumentFields = map[string]bool{ + "id": true, "document_id": true, "collection_id": true, "name": true, + "description": true, "status": true, "type": true, "size": true, + "segment_count": true, "job_id": true, "uploader_id": true, "tags": true, + "locale": true, "system": true, "readonly": true, "sort": true, "cover": true, + "file_id": true, "file_name": true, "file_mime_type": true, + "url": true, "url_title": true, "text_content": true, + "converter_provider_id": true, "converter_option_id": true, "converter_properties": true, + "fetcher_provider_id": true, "fetcher_option_id": true, "fetcher_properties": true, + "chunking_provider_id": true, "chunking_option_id": true, "chunking_properties": true, + "extraction_provider_id": true, "extraction_option_id": true, "extraction_properties": true, + "processed_at": true, "error_message": true, "created_at": true, "updated_at": true, + } + + // DefaultDocumentFields defines the default compact field list + DefaultDocumentFields = []interface{}{ + "id", "document_id", "collection_id", "name", "description", + "cover", "tags", "type", "size", "segment_count", "status", "locale", + "system", "readonly", "file_id", "file_name", "file_mime_type", "uploader_id", + "url", "url_title", "text_content", "job_id", + "error_message", "created_at", "updated_at", + } + + // ValidDocumentSortFields defines valid fields for sorting + ValidDocumentSortFields = map[string]bool{ + "created_at": true, + "updated_at": true, + "name": true, + "size": true, + "segment_count": true, + "sort": true, + "processed_at": true, + } +) + +// DefaultDocumentSort defines the default sort order for document queries +var DefaultDocumentSort = []model.QueryOrder{ + {Column: DefaultSortField, Option: DefaultSortOrder}, +} + +// Default uploader +const ( + DefaultUploader = "local" +) diff --git a/kb/api/document.go b/kb/api/document.go new file mode 100644 index 00000000..faa6d3aa --- /dev/null +++ b/kb/api/document.go @@ -0,0 +1,279 @@ +package api + +import ( + "context" + "fmt" + + "github.com/yaoapp/gou/model" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/maps" +) + +// ListDocuments lists documents with pagination and filtering +func (instance *KBInstance) ListDocuments(ctx context.Context, filter *ListDocumentsFilter) (*ListDocumentsResult, 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 = DefaultDocumentFields + } else { + // Filter valid fields + validFields := []interface{}{} + for _, field := range selectFields { + if fieldStr, ok := field.(string); ok && AvailableDocumentFields[fieldStr] { + validFields = append(validFields, field) + } + } + if len(validFields) == 0 { + selectFields = DefaultDocumentFields + } 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 collection_id + if filter.CollectionID != "" { + wheres = append(wheres, model.QueryWhere{ + Column: "collection_id", + Value: filter.CollectionID, + }) + } + + // 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 tag + if filter.Tag != "" { + wheres = append(wheres, model.QueryWhere{ + Column: "tags", + Value: "%" + filter.Tag + "%", + OP: "like", + }) + } + + // 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 status_not (exclude specific statuses) + if len(filter.StatusNot) > 0 { + for _, status := range filter.StatusNot { + if status != "" { + wheres = append(wheres, model.QueryWhere{ + Column: "status", + Value: status, + OP: "!=", + }) + } + } + } + + param.Wheres = wheres + + // Process sort orders + orders := filter.Sort + if len(orders) == 0 { + orders = DefaultDocumentSort + } else { + // Validate sort fields + validOrders := []model.QueryOrder{} + for _, order := range orders { + if ValidDocumentSortFields[order.Column] { + // Validate sort order + if order.Option != "asc" && order.Option != "desc" { + order.Option = "desc" + } + validOrders = append(validOrders, order) + } + } + if len(validOrders) == 0 { + orders = DefaultDocumentSort + } else { + orders = validOrders + } + } + + param.Orders = orders + + // Query documents + result, err := instance.Config.SearchDocuments(param, page, pageSize) + if err != nil { + return nil, fmt.Errorf("failed to search documents: %w", err) + } + + // Convert result to ListDocumentsResult + listResult := &ListDocumentsResult{ + Page: page, + PageSize: pageSize, + Data: make([]map[string]interface{}, 0), + } + + // 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 { + 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 { + 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 +} + +// GetDocument retrieves a document by ID +func (instance *KBInstance) GetDocument(ctx context.Context, docID string, params *GetDocumentParams) (map[string]interface{}, error) { + if docID == "" { + return nil, fmt.Errorf("document ID is required") + } + + // Process select fields + var selectFields []interface{} + if params != nil && len(params.Select) > 0 { + for _, field := range params.Select { + if fieldStr, ok := field.(string); ok && AvailableDocumentFields[fieldStr] { + selectFields = append(selectFields, field) + } + } + } + if len(selectFields) == 0 { + selectFields = DefaultDocumentFields + } + + // Build query parameters + param := model.QueryParam{ + Select: selectFields, + } + + // Query single document + result, err := instance.Config.FindDocument(docID, param) + if err != nil { + return nil, err + } + + return result, nil +} + +// RemoveDocuments removes documents by IDs +func (instance *KBInstance) RemoveDocuments(ctx context.Context, params *RemoveDocumentsParams) (*RemoveDocumentsResult, error) { + if len(params.DocumentIDs) == 0 { + return nil, fmt.Errorf("document IDs are required") + } + + // Remove documents using GraphRag + deletedCount, err := instance.GraphRag.RemoveDocs(ctx, params.DocumentIDs) + if err != nil { + return nil, fmt.Errorf("failed to remove documents from GraphRag: %w", err) + } + + // Also remove documents from the database and track collections to update + dbDeletedCount := 0 + collectionsToUpdate := make(map[string]bool) + + for _, docID := range params.DocumentIDs { + // Get document info before deletion to track collection + if docInfo, err := instance.Config.FindDocument(docID, model.QueryParam{ + Select: []interface{}{"collection_id"}, + }); err == nil && docInfo != nil { + if collectionID, ok := docInfo["collection_id"].(string); ok && collectionID != "" { + collectionsToUpdate[collectionID] = true + } + } + + if err := instance.Config.RemoveDocument(docID); err != nil { + return nil, fmt.Errorf("failed to remove document from database: %w", err) + } + dbDeletedCount++ + } + + // Update document counts for affected collections and sync to GraphRag + for collectionID := range collectionsToUpdate { + if err := instance.updateDocumentCountWithSync(ctx, collectionID); err != nil { + log.Error("Failed to update document count for collection %s: %v", collectionID, err) + } + } + + return &RemoveDocumentsResult{ + Message: "Documents removed successfully", + DeletedCount: deletedCount, + RequestedCount: len(params.DocumentIDs), + DBDeletedCount: dbDeletedCount, + }, nil +} diff --git a/kb/api/document_test.go b/kb/api/document_test.go new file mode 100644 index 00000000..44a9b6be --- /dev/null +++ b/kb/api/document_test.go @@ -0,0 +1,291 @@ +package api_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go, which handles environment setup +// Run tests with: source env.local.sh && go test -v ./kb/api/... + +// createTestCollectionForDoc is a helper to create a test collection for document tests +func createTestCollectionForDoc(t *testing.T, ctx context.Context) string { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + collectionID := fmt.Sprintf("test_doc_%d", time.Now().UnixNano()) + + params := &api.CreateCollectionParams{ + ID: collectionID, + Metadata: map[string]interface{}{ + "name": "Test Document Collection", + "description": "Collection for document tests", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + + _, err := kb.API.CreateCollection(ctx, params) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + + return collectionID +} + +// cleanupTestCollectionForDoc removes a test collection +func cleanupTestCollectionForDoc(ctx context.Context, collectionID string) { + if kb.API != nil { + _, _ = kb.API.RemoveCollection(ctx, collectionID) + } +} + +// addTestDocument adds a test document and returns its ID +func addTestDocument(t *testing.T, ctx context.Context, collectionID, title string) string { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: fmt.Sprintf("Test document content for %s", title), + Metadata: map[string]interface{}{ + "title": title, + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + } + result, err := kb.API.AddText(ctx, params) + if err != nil { + t.Fatalf("Failed to add test document: %v", err) + } + return result.DocID +} + +// ========== ListDocuments Tests ========== + +func TestListDocuments(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForDoc(t, ctx) + defer cleanupTestCollectionForDoc(ctx, collectionID) + + // Add some test documents + for i := 0; i < 3; i++ { + addTestDocument(t, ctx, collectionID, fmt.Sprintf("Test Document %d", i+1)) + } + + t.Run("ListDocumentsDefault", func(t *testing.T) { + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + } + + result, err := kb.API.ListDocuments(ctx, filter) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.GreaterOrEqual(t, len(result.Data), 3) + assert.Equal(t, 1, result.Page) + assert.Equal(t, 20, result.PageSize) + t.Logf("Found %d documents in collection", len(result.Data)) + }) + + t.Run("ListDocumentsWithPagination", func(t *testing.T) { + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 2, + CollectionID: collectionID, + } + + result, err := kb.API.ListDocuments(ctx, filter) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.LessOrEqual(t, len(result.Data), 2) + }) + + t.Run("ListDocumentsWithKeywords", func(t *testing.T) { + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + Keywords: "Test Document 1", + } + + result, err := kb.API.ListDocuments(ctx, filter) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.GreaterOrEqual(t, len(result.Data), 1) + }) + + t.Run("ListDocumentsWithStatus", func(t *testing.T) { + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + Status: []string{"completed"}, + } + + result, err := kb.API.ListDocuments(ctx, filter) + assert.NoError(t, err) + assert.NotNil(t, result) + + for _, doc := range result.Data { + status, ok := doc["status"].(string) + if ok { + assert.Equal(t, "completed", status) + } + } + }) + + t.Run("ListDocumentsEmptyResult", func(t *testing.T) { + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: collectionID, + Keywords: "nonexistent_keyword_xyz123", + } + + result, err := kb.API.ListDocuments(ctx, filter) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Equal(t, 0, len(result.Data)) + }) +} + +// ========== GetDocument Tests ========== + +func TestGetDocument(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForDoc(t, ctx) + defer cleanupTestCollectionForDoc(ctx, collectionID) + + // Add a test document + docID := addTestDocument(t, ctx, collectionID, "GetDocument Test") + + t.Run("GetDocumentSuccess", func(t *testing.T) { + doc, err := kb.API.GetDocument(ctx, docID, nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.Equal(t, docID, doc["document_id"]) + assert.Equal(t, collectionID, doc["collection_id"]) + assert.Equal(t, "GetDocument Test", doc["name"]) + assert.Equal(t, "text", doc["type"]) + t.Logf("Retrieved document: %v", doc["name"]) + }) + + t.Run("GetDocumentWithSelect", func(t *testing.T) { + params := &api.GetDocumentParams{ + Select: []interface{}{"document_id", "name", "type", "status"}, + } + + doc, err := kb.API.GetDocument(ctx, docID, params) + assert.NoError(t, err) + assert.NotNil(t, doc) + assert.NotNil(t, doc["document_id"]) + assert.NotNil(t, doc["name"]) + }) + + t.Run("GetDocumentNotFound", func(t *testing.T) { + doc, err := kb.API.GetDocument(ctx, "nonexistent_doc_id", nil) + assert.Error(t, err) + assert.Nil(t, doc) + assert.Contains(t, err.Error(), "not found") + }) + + t.Run("GetDocumentEmptyID", func(t *testing.T) { + doc, err := kb.API.GetDocument(ctx, "", nil) + assert.Error(t, err) + assert.Nil(t, doc) + assert.Contains(t, err.Error(), "required") + }) +} + +// ========== RemoveDocuments Tests ========== + +func TestRemoveDocuments(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + collectionID := createTestCollectionForDoc(t, ctx) + defer cleanupTestCollectionForDoc(ctx, collectionID) + + // Add test documents + var docIDs []string + for i := 0; i < 3; i++ { + docID := addTestDocument(t, ctx, collectionID, fmt.Sprintf("Remove Test %d", i+1)) + docIDs = append(docIDs, docID) + } + + t.Run("RemoveDocumentsSuccess", func(t *testing.T) { + params := &api.RemoveDocumentsParams{ + DocumentIDs: docIDs[:2], // Remove first 2 documents + } + + result, err := kb.API.RemoveDocuments(ctx, params) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Equal(t, 2, result.RequestedCount) + assert.GreaterOrEqual(t, result.DeletedCount, 0) + t.Logf("Removed documents: requested=%d, deleted=%d", result.RequestedCount, result.DeletedCount) + + // Verify documents are removed + for _, docID := range docIDs[:2] { + doc, err := kb.API.GetDocument(ctx, docID, nil) + assert.Error(t, err) + assert.Nil(t, doc) + } + + // Verify remaining document still exists + doc, err := kb.API.GetDocument(ctx, docIDs[2], nil) + assert.NoError(t, err) + assert.NotNil(t, doc) + }) + + t.Run("RemoveDocumentsEmptyList", func(t *testing.T) { + params := &api.RemoveDocumentsParams{ + DocumentIDs: []string{}, + } + + result, err := kb.API.RemoveDocuments(ctx, params) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "required") + }) + + t.Run("RemoveDocumentsNonexistent", func(t *testing.T) { + params := &api.RemoveDocumentsParams{ + DocumentIDs: []string{"nonexistent_doc_1", "nonexistent_doc_2"}, + } + + result, err := kb.API.RemoveDocuments(ctx, params) + // Should succeed but with 0 deleted + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Equal(t, 2, result.RequestedCount) + }) +} diff --git a/kb/api/interfaces.go b/kb/api/interfaces.go index 1436341c..5ac527b7 100644 --- a/kb/api/interfaces.go +++ b/kb/api/interfaces.go @@ -17,10 +17,20 @@ type API interface { 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) - // ... + // Document operations + ListDocuments(ctx context.Context, filter *ListDocumentsFilter) (*ListDocumentsResult, error) + GetDocument(ctx context.Context, docID string, params *GetDocumentParams) (map[string]interface{}, error) + RemoveDocuments(ctx context.Context, params *RemoveDocumentsParams) (*RemoveDocumentsResult, error) + + // Document add operations (sync) + AddFile(ctx context.Context, params *AddFileParams) (*AddDocumentResult, error) + AddText(ctx context.Context, params *AddTextParams) (*AddDocumentResult, error) + AddURL(ctx context.Context, params *AddURLParams) (*AddDocumentResult, error) + + // Document add operations (async) + AddFileAsync(ctx context.Context, params *AddFileParams) (*AddDocumentAsyncResult, error) + AddTextAsync(ctx context.Context, params *AddTextParams) (*AddDocumentAsyncResult, error) + AddURLAsync(ctx context.Context, params *AddURLParams) (*AddDocumentAsyncResult, error) // Segment operations (future) // ... @@ -28,7 +38,8 @@ type API interface { // KBInstance holds the KB instance dependencies required by the API type KBInstance struct { - GraphRag types.GraphRag // GraphRag instance for vector/graph operations + 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 index 8a9b6af0..a2997b26 100644 --- a/kb/api/types.go +++ b/kb/api/types.go @@ -71,3 +71,127 @@ type UpdateMetadataResult struct { CollectionID string `json:"collection_id" yaml:"collection_id"` Message string `json:"message" yaml:"message"` } + +// ========== Document Types ========== + +// ListDocumentsFilter represents the filter options for listing documents +type ListDocumentsFilter struct { + Page int `json:"page" yaml:"page"` + PageSize int `json:"pagesize" yaml:"pagesize"` + CollectionID string `json:"collection_id,omitempty" yaml:"collection_id,omitempty"` + Keywords string `json:"keywords,omitempty" yaml:"keywords,omitempty"` + Tag string `json:"tag,omitempty" yaml:"tag,omitempty"` + Status []string `json:"status,omitempty" yaml:"status,omitempty"` + StatusNot []string `json:"status_not,omitempty" yaml:"status_not,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 +} + +// ListDocumentsResult represents the result of listing documents +type ListDocumentsResult 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"` +} + +// GetDocumentParams represents the parameters for getting a document +type GetDocumentParams struct { + Select []interface{} `json:"select,omitempty" yaml:"select,omitempty"` +} + +// RemoveDocumentsParams represents the parameters for removing documents +type RemoveDocumentsParams struct { + DocumentIDs []string `json:"document_ids" yaml:"document_ids"` +} + +// RemoveDocumentsResult represents the result of removing documents +type RemoveDocumentsResult struct { + Message string `json:"message" yaml:"message"` + DeletedCount int `json:"deleted_count" yaml:"deleted_count"` + RequestedCount int `json:"requested_count" yaml:"requested_count"` + DBDeletedCount int `json:"db_deleted_count" yaml:"db_deleted_count"` +} + +// AddFileParams represents the parameters for adding a file +type AddFileParams struct { + CollectionID string `json:"collection_id" yaml:"collection_id"` + FileID string `json:"file_id" yaml:"file_id"` + Uploader string `json:"uploader,omitempty" yaml:"uploader,omitempty"` + DocID string `json:"doc_id,omitempty" yaml:"doc_id,omitempty"` + Locale string `json:"locale,omitempty" yaml:"locale,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` + Chunking *ProviderConfigParams `json:"chunking" yaml:"chunking"` + Embedding *ProviderConfigParams `json:"embedding" yaml:"embedding"` + Extraction *ProviderConfigParams `json:"extraction,omitempty" yaml:"extraction,omitempty"` + Fetcher *ProviderConfigParams `json:"fetcher,omitempty" yaml:"fetcher,omitempty"` + Converter *ProviderConfigParams `json:"converter,omitempty" yaml:"converter,omitempty"` + Job *JobOptionsParams `json:"job,omitempty" yaml:"job,omitempty"` + AuthScope map[string]interface{} `json:"-" yaml:"-"` // Internal: authentication scope fields +} + +// AddTextParams represents the parameters for adding text +type AddTextParams struct { + CollectionID string `json:"collection_id" yaml:"collection_id"` + Text string `json:"text" yaml:"text"` + DocID string `json:"doc_id,omitempty" yaml:"doc_id,omitempty"` + Locale string `json:"locale,omitempty" yaml:"locale,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` + Chunking *ProviderConfigParams `json:"chunking" yaml:"chunking"` + Embedding *ProviderConfigParams `json:"embedding" yaml:"embedding"` + Extraction *ProviderConfigParams `json:"extraction,omitempty" yaml:"extraction,omitempty"` + Fetcher *ProviderConfigParams `json:"fetcher,omitempty" yaml:"fetcher,omitempty"` + Converter *ProviderConfigParams `json:"converter,omitempty" yaml:"converter,omitempty"` + Job *JobOptionsParams `json:"job,omitempty" yaml:"job,omitempty"` + AuthScope map[string]interface{} `json:"-" yaml:"-"` // Internal: authentication scope fields +} + +// AddURLParams represents the parameters for adding a URL +type AddURLParams struct { + CollectionID string `json:"collection_id" yaml:"collection_id"` + URL string `json:"url" yaml:"url"` + DocID string `json:"doc_id,omitempty" yaml:"doc_id,omitempty"` + Locale string `json:"locale,omitempty" yaml:"locale,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` + Chunking *ProviderConfigParams `json:"chunking" yaml:"chunking"` + Embedding *ProviderConfigParams `json:"embedding" yaml:"embedding"` + Extraction *ProviderConfigParams `json:"extraction,omitempty" yaml:"extraction,omitempty"` + Fetcher *ProviderConfigParams `json:"fetcher,omitempty" yaml:"fetcher,omitempty"` + Converter *ProviderConfigParams `json:"converter,omitempty" yaml:"converter,omitempty"` + Job *JobOptionsParams `json:"job,omitempty" yaml:"job,omitempty"` + AuthScope map[string]interface{} `json:"-" yaml:"-"` // Internal: authentication scope fields +} + +// ProviderConfigParams represents a provider configuration +type ProviderConfigParams struct { + ProviderID string `json:"provider_id" yaml:"provider_id"` + OptionID string `json:"option_id,omitempty" yaml:"option_id,omitempty"` + Properties map[string]interface{} `json:"properties,omitempty" yaml:"properties,omitempty"` +} + +// JobOptionsParams contains job options for async operations +type JobOptionsParams struct { + Name string `json:"name,omitempty" yaml:"name,omitempty"` + Description string `json:"description,omitempty" yaml:"description,omitempty"` + Icon string `json:"icon,omitempty" yaml:"icon,omitempty"` + Category string `json:"category,omitempty" yaml:"category,omitempty"` +} + +// AddDocumentResult represents the result of adding a document (sync) +type AddDocumentResult struct { + Message string `json:"message" yaml:"message"` + CollectionID string `json:"collection_id" yaml:"collection_id"` + DocID string `json:"doc_id" yaml:"doc_id"` + FileID string `json:"file_id,omitempty" yaml:"file_id,omitempty"` + URL string `json:"url,omitempty" yaml:"url,omitempty"` +} + +// AddDocumentAsyncResult represents the result of adding a document (async) +type AddDocumentAsyncResult struct { + JobID string `json:"job_id" yaml:"job_id"` + DocID string `json:"doc_id" yaml:"doc_id"` +} diff --git a/kb/api/utils.go b/kb/api/utils.go new file mode 100644 index 00000000..169dc166 --- /dev/null +++ b/kb/api/utils.go @@ -0,0 +1,325 @@ +package api + +import ( + "context" + "fmt" + + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/maps" + "github.com/yaoapp/yao/kb/providers/factory" + kbtypes "github.com/yaoapp/yao/kb/types" +) + +// updateDocumentAfterProcessing updates document status and segment count after processing +func (instance *KBInstance) updateDocumentAfterProcessing(ctx context.Context, docID, collectionID string) { + // Update status to completed + if err := instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "completed"}); err != nil { + log.Error("Failed to update document status to completed: %v", err) + } + + // Update segment count + if segmentCount, err := instance.GraphRag.SegmentCount(ctx, docID); err != nil { + log.Error("Failed to get segment count for document %s: %v", docID, err) + } else { + if err := instance.Config.UpdateSegmentCount(docID, segmentCount); err != nil { + log.Error("Failed to update segment count for document %s: %v", docID, err) + } + } + + // Update document count for collection + if err := instance.updateDocumentCountWithSync(ctx, collectionID); err != nil { + log.Error("Failed to update document count for collection %s: %v", collectionID, err) + } +} + +// updateDocumentCountWithSync updates document count and syncs to GraphRag +func (instance *KBInstance) updateDocumentCountWithSync(ctx context.Context, collectionID string) error { + // Get document count + count, err := instance.Config.DocumentCount(collectionID) + if err != nil { + return fmt.Errorf("failed to get document count: %w", err) + } + + // Update collection in database + if err := instance.Config.UpdateCollection(collectionID, maps.MapStrAny{"document_count": count}); err != nil { + return fmt.Errorf("failed to update collection document count: %w", err) + } + + // Sync to GraphRag + metadata := map[string]interface{}{"document_count": count} + if err := instance.GraphRag.UpdateCollectionMetadata(ctx, collectionID, metadata); err != nil { + return fmt.Errorf("failed to sync document count to GraphRag: %w", err) + } + + return nil +} + +// toUpsertOptions converts provider config params to UpsertOptions +func (instance *KBInstance) toUpsertOptions(docID, collectionID, locale, filename, contentType string, chunking, embedding, extraction, fetcher, converter *ProviderConfigParams) (*graphragtypes.UpsertOptions, error) { + if locale == "" { + locale = DefaultLocale + } + + options := &graphragtypes.UpsertOptions{ + CollectionID: collectionID, + DocID: docID, + } + + // Create chunking provider + if chunking != nil { + chunkingOption, err := instance.getProviderOption("chunking", chunking.ProviderID, chunking.OptionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve chunking provider: %w", err) + } + + chunkingProvider, err := factory.MakeChunking(chunking.ProviderID, chunkingOption) + if err != nil { + return nil, fmt.Errorf("failed to create chunking provider: %w", err) + } + options.Chunking = chunkingProvider + + chunkingOpts, err := factory.ChunkingOptions(chunking.ProviderID, chunkingOption) + if err != nil { + return nil, fmt.Errorf("failed to get chunking options: %w", err) + } + options.ChunkingOptions = chunkingOpts + } + + // Create embedding provider + if embedding != nil { + embeddingOption, err := instance.getProviderOption("embedding", embedding.ProviderID, embedding.OptionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve embedding provider: %w", err) + } + + embeddingProvider, err := factory.MakeEmbedding(embedding.ProviderID, embeddingOption) + if err != nil { + return nil, fmt.Errorf("failed to create embedding provider: %w", err) + } + options.Embedding = embeddingProvider + } + + // Create extraction provider (optional, but required if graph is enabled) + // If extraction is not provided, try to use the default extraction provider + if extraction != nil { + extractionOption, err := instance.getProviderOption("extraction", extraction.ProviderID, extraction.OptionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve extraction provider: %w", err) + } + + extractionProvider, err := factory.MakeExtraction(extraction.ProviderID, extractionOption) + if err != nil { + return nil, fmt.Errorf("failed to create extraction provider: %w", err) + } + options.Extraction = extractionProvider + } else { + // Try to get default extraction provider to avoid gou's DetectExtractor with hardcoded connector + defaultExtraction := instance.getDefaultProvider("extraction", locale) + if defaultExtraction != nil { + extractionOption, err := instance.getProviderOption("extraction", defaultExtraction.ID, "", locale) + if err == nil { + extractionProvider, err := factory.MakeExtraction(defaultExtraction.ID, extractionOption) + if err == nil { + options.Extraction = extractionProvider + } + } + } + } + + // Create fetcher provider (optional) + if fetcher != nil { + fetcherOption, err := instance.getProviderOption("fetcher", fetcher.ProviderID, fetcher.OptionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve fetcher provider: %w", err) + } + + fetcherProvider, err := factory.MakeFetcher(fetcher.ProviderID, fetcherOption) + if err != nil { + return nil, fmt.Errorf("failed to create fetcher provider: %w", err) + } + options.Fetcher = fetcherProvider + } + + // Create converter provider (optional or auto-detect) + if converter != nil { + converterOption, err := instance.getProviderOption("converter", converter.ProviderID, converter.OptionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve converter provider: %w", err) + } + + converterProvider, err := factory.MakeConverter(converter.ProviderID, converterOption) + if err != nil { + return nil, fmt.Errorf("failed to create converter provider: %w", err) + } + options.Converter = converterProvider + } else if filename != "" || contentType != "" { + // Auto-detect converter + matched, converterID, err := factory.AutoDetectConverter(filename, contentType) + if err != nil { + return nil, fmt.Errorf("failed to auto-detect converter: %w", err) + } + + if matched { + converterOption, err := instance.getProviderOption("converter", converterID, "", locale) + if err != nil { + return nil, fmt.Errorf("failed to resolve auto-detected converter provider: %w", err) + } + + converterProvider, err := factory.MakeConverter(converterID, converterOption) + if err != nil { + return nil, fmt.Errorf("failed to create auto-detected converter provider: %w", err) + } + options.Converter = converterProvider + } + } + + return options, nil +} + +// getProviderOption gets a provider option by provider type, ID and option ID +func (instance *KBInstance) getProviderOption(providerType, providerID, optionID, locale string) (*kbtypes.ProviderOption, error) { + provider, err := instance.Providers.GetProvider(providerType, providerID, locale) + if err != nil { + return nil, fmt.Errorf("provider %s not found for locale %s: %w", providerID, locale, err) + } + + if optionID != "" { + option, exists := provider.GetOption(optionID) + if !exists { + return nil, fmt.Errorf("option %s not found in provider %s", optionID, providerID) + } + return option, nil + } + + // Return default option + if provider.Options != nil { + for _, option := range provider.Options { + if option.Default { + return option, nil + } + } + if len(provider.Options) > 0 { + return provider.Options[0], nil + } + } + + return nil, fmt.Errorf("no option specified and no default option found for provider %s", providerID) +} + +// getDefaultProvider returns the default provider for a given type and locale +func (instance *KBInstance) getDefaultProvider(providerType, locale string) *kbtypes.Provider { + if instance.Providers == nil { + return nil + } + + providers := instance.Providers.GetProviders(providerType, locale) + if len(providers) == 0 { + return nil + } + + // Find provider with default=true + for _, provider := range providers { + if provider.Default { + return provider + } + } + + // Return first provider if no default is set + return providers[0] +} + +// getJobOptions returns job options with defaults +func getJobOptions(job *JobOptionsParams, defaultName, defaultDescription, defaultIcon, defaultCategory string) (string, string, string, string) { + name := defaultName + description := defaultDescription + icon := defaultIcon + category := defaultCategory + + if job != nil { + if job.Name != "" { + name = job.Name + } + if job.Description != "" { + description = job.Description + } + if job.Icon != "" { + icon = job.Icon + } + if job.Category != "" { + category = job.Category + } + } + + return name, description, icon, category +} + +// addBaseFieldsFromParams adds base fields from parameters to document data +func addBaseFieldsFromParams(data map[string]interface{}, locale string, metadata map[string]interface{}, chunking, embedding, extraction, fetcher, converter *ProviderConfigParams) { + if locale != "" { + data["locale"] = locale + } + + // Extract fields from metadata + if metadata != nil { + if description, ok := metadata["description"]; ok && description != nil { + data["description"] = description + } + if cover, ok := metadata["cover"]; ok && cover != nil { + data["cover"] = cover + } + if tags, ok := metadata["tags"]; ok && tags != nil { + data["tags"] = tags + } + if name, ok := metadata["name"]; ok && name != nil { + data["name"] = name + } + } + + // Add provider configurations + if converter != nil { + data["converter_provider_id"] = converter.ProviderID + if converter.OptionID != "" { + data["converter_option_id"] = converter.OptionID + } + if converter.Properties != nil { + data["converter_properties"] = converter.Properties + } + } + if fetcher != nil { + data["fetcher_provider_id"] = fetcher.ProviderID + if fetcher.OptionID != "" { + data["fetcher_option_id"] = fetcher.OptionID + } + if fetcher.Properties != nil { + data["fetcher_properties"] = fetcher.Properties + } + } + if chunking != nil { + data["chunking_provider_id"] = chunking.ProviderID + if chunking.OptionID != "" { + data["chunking_option_id"] = chunking.OptionID + } + if chunking.Properties != nil { + data["chunking_properties"] = chunking.Properties + } + } + if embedding != nil { + data["embedding_provider_id"] = embedding.ProviderID + if embedding.OptionID != "" { + data["embedding_option_id"] = embedding.OptionID + } + if embedding.Properties != nil { + data["embedding_properties"] = embedding.Properties + } + } + if extraction != nil { + data["extraction_provider_id"] = extraction.ProviderID + if extraction.OptionID != "" { + data["extraction_option_id"] = extraction.OptionID + } + if extraction.Properties != nil { + data["extraction_properties"] = extraction.Properties + } + } +} diff --git a/openapi/kb/addfile.go b/openapi/kb/addfile.go index 0109e4f1..f1d90cdf 100644 --- a/openapi/kb/addfile.go +++ b/openapi/kb/addfile.go @@ -6,239 +6,24 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/utils" - "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/attachment" - "github.com/yaoapp/yao/job" "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" ) -// CreateDocumentRecord creates a document record in the database immediately -// This is called synchronously when the API request comes in -func CreateDocumentRecord(ctx context.Context, authInfo *oauthtypes.AuthorizedInfo, req *AddFileRequest, jobID string) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get file manager - m, ok := attachment.Managers[req.Uploader] - if !ok { - return fmt.Errorf("invalid uploader: %s not found", req.Uploader) - } - - // Check if the file exists - exists := m.Exists(ctx, req.FileID) - if !exists { - return fmt.Errorf("file not found: %s", req.FileID) - } - - // Get file info and path - path, contentType, err := m.LocalPath(ctx, req.FileID) - if err != nil { - return fmt.Errorf("failed to get local path: %w", err) - } - - fileInfo, err := m.Info(ctx, req.FileID) - if err != nil { - return fmt.Errorf("failed to get file info: %w", err) - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Prepare document data for database - documentData := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": fileInfo.Filename, - "type": "file", - "status": "pending", - "uploader_id": req.Uploader, - "file_id": req.FileID, - "file_name": fileInfo.Filename, - "file_path": path, - "file_mime_type": contentType, - "size": int64(fileInfo.Bytes), - "job_id": jobID, - } - - // With create scope - if authInfo != nil { - documentData = authInfo.WithCreateScope(documentData) - } - - // Add base request fields - req.BaseUpsertRequest.AddBaseFields(documentData) - - // Create database record - _, err = config.CreateDocument(maps.MapStrAny(documentData)) - if err != nil { - return fmt.Errorf("failed to save document metadata: %w", err) - } - - return nil -} - -// HandleFileContent processes the actual file content and updates the knowledge base -// This is called asynchronously by the job system -func HandleFileContent(ctx context.Context, req *AddFileRequest) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get file manager - m, ok := attachment.Managers[req.Uploader] - if !ok { - return fmt.Errorf("invalid uploader: %s not found", req.Uploader) - } - - // Get file info and path - path, contentType, err := m.LocalPath(ctx, req.FileID) - if err != nil { - return fmt.Errorf("failed to get local path: %w", err) - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Convert request to UpsertOptions - upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions(path, contentType) - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to convert request to upsert options: %w", err) - } - - // Perform upsert operation with file path - _, err = kb.Instance.AddFile(ctx, path, upsertOptions) - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to add file: %w", err) - } - - // Update status to completed after successful processing - if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { - log.Error("Failed to update document status to completed: %v", err) - } - - // Update segment count for the document - if segmentCount, err := kb.Instance.SegmentCount(ctx, req.DocID); err != nil { - log.Error("Failed to get segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Got segment count %d for document %s", segmentCount, req.DocID) - if err := config.UpdateSegmentCount(req.DocID, segmentCount); err != nil { - log.Error("Failed to update segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Successfully updated segment count to %d for document %s", segmentCount, req.DocID) - } - } - - // Update document count for the collection and sync to GraphRag - if err := UpdateDocumentCountWithSync(req.CollectionID, config); err != nil { - log.Error("Failed to update document count for collection %s: %v", req.CollectionID, err) - } else { - log.Info("Successfully updated document count for collection %s", req.CollectionID) - } - - return nil -} - -// AddFileHandler processes a file addition request with business logic only -// This function combines both document creation and content processing for sync operations -func AddFileHandler(ctx context.Context, authInfo *oauthtypes.AuthorizedInfo, req *AddFileRequest, jobID ...string) error { - // Validate request - if err := req.Validate(); err != nil { - return err - } - - // DocID should be generated by the caller before calling this function - if req.DocID == "" { - return fmt.Errorf("document ID is required") - } - - // For sync operations, create document record and process content immediately - var jid string - if len(jobID) > 0 { - jid = jobID[0] - } - - // Create document record - if err := CreateDocumentRecord(ctx, authInfo, req, jid); err != nil { - return err - } - - // Process file content - return HandleFileContent(ctx, req) -} - -// addFileWithRequest processes a file addition with pre-parsed request using Gin context -func addFileWithRequest(c *gin.Context, req *AddFileRequest) { - - // Check collection permission - authInfo := authorized.GetInfo(c) - hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: err.Error(), - } - response.RespondWithError(c, response.StatusForbidden, errorResp) - return - } - - // 403 Forbidden - if !hasPermission { - errorResp := &response.ErrorResponse{ - Code: response.ErrAccessDenied.Code, - ErrorDescription: "Forbidden: No permission to update collection", - } - response.RespondWithError(c, response.StatusForbidden, errorResp) - return - } - - // Use the business logic function - err = AddFileHandler(c.Request.Context(), authInfo, req) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Return success response - result := gin.H{ - "message": "File added successfully", - "collection_id": req.CollectionID, - "file_id": req.FileID, - "doc_id": req.DocID, - } - - response.RespondWithSuccess(c, response.StatusCreated, result) -} - -// AddFile adds a file to a collection +// AddFile adds a file to a collection (sync) func AddFile(c *gin.Context) { var req AddFileRequest - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -267,8 +52,44 @@ func AddFile(c *gin.Context) { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } - // Process the request - addFileWithRequest(c, &req) + // Check collection permission + authInfo := authorized.GetInfo(c) + hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // 403 Forbidden + if !hasPermission { + errorResp := &response.ErrorResponse{ + Code: response.ErrAccessDenied.Code, + ErrorDescription: "Forbidden: No permission to update collection", + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // Convert request to API params + params := convertAddFileRequest(&req, authInfo) + + // Call kb.API + result, err := kb.API.AddFile(c.Request.Context(), params) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Return success response + response.RespondWithSuccess(c, response.StatusCreated, result) } // AddFileAsync adds file to a collection asynchronously @@ -277,9 +98,9 @@ func AddFileAsync(c *gin.Context) { log.Info("AddFileAsync: Starting async file addition") - // Check if kb.Instance is available - if !checkKBInstance(c) { - log.Error("AddFileAsync: KB instance check failed") + // Check if kb.API is available + if !checkKBAPI(c) { + log.Error("AddFileAsync: KB API check failed") return } @@ -294,7 +115,7 @@ func AddFileAsync(c *gin.Context) { return } - log.Info("AddFileAsync: Request parsed successfully: %+v", req) + log.Info("AddFileAsync: Request parsed successfully") // Validate request if err := req.Validate(); err != nil { @@ -309,24 +130,14 @@ func AddFileAsync(c *gin.Context) { log.Info("AddFileAsync: Request validation passed") - // Validate file and get path - _, _, err := validateFileAndGetPath(c, &req) - if err != nil { + // Validate file exists + if err := validateFileExists(c, &req); err != nil { log.Error("AddFileAsync: File validation failed: %v", err) return } log.Info("AddFileAsync: File validation passed") - // Convert request to UpsertOptions (just for validation) - _, err = getUpsertOptions(c, &req.BaseUpsertRequest) - if err != nil { - log.Error("AddFileAsync: UpsertOptions validation failed: %v", err) - return - } - - log.Info("AddFileAsync: UpsertOptions validation passed") - // Generate document ID if not provided if req.DocID == "" { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) @@ -356,112 +167,25 @@ func AddFileAsync(c *gin.Context) { return } - // Step 1: Get job options with defaults - jobName, jobDescription, jobIcon, jobCategory := req.GetJobOptions( - "Knowledge Base File Processing", // default name - "Processing and indexing file content for knowledge base search", // default description - "library_add", // default icon (Material Icon) - "Knowledge Base", // default category - ) + // Convert request to API params + params := convertAddFileRequest(&req, authInfo) - // Create job data - jobCreateData := map[string]interface{}{ - "name": jobName, - "description": jobDescription, - "category_name": jobCategory, // Pass category name directly, let SaveJob handle it - } - if jobIcon != "" { - jobCreateData["icon"] = jobIcon - } - - // With create scope - if authInfo != nil { - jobCreateData = authInfo.WithCreateScope(jobCreateData) - } - - // Create and save Job in one step to get JobID - j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + // Call kb.API async + result, err := kb.API.AddFileAsync(c.Request.Context(), params) if err != nil { - log.Error("AddFileAsync: Job creation and save failed: %v", err) + log.Error("AddFileAsync: Failed to add file async: %v", err) errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create and save job: " + err.Error(), + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } - log.Info("AddFileAsync: Job created and saved with ID: %s", j.JobID) - - // Step 2: Create document record immediately with job_id - err = CreateDocumentRecord(c.Request.Context(), authInfo, &req, j.JobID) - if err != nil { - log.Error("AddFileAsync: Failed to create document record: %v", err) - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create document record: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddFileAsync: Document record created successfully") - - // Step 4: Add execution to job - jobData := map[string]interface{}{ - "collection_id": req.CollectionID, - "file_id": req.FileID, - "uploader": req.Uploader, - "locale": req.Locale, - "doc_id": req.DocID, - "metadata": req.Metadata, - "chunking": req.Chunking, - "embedding": req.Embedding, - "extraction": req.Extraction, - "fetcher": req.Fetcher, - "converter": req.Converter, - } - - err = j.Add(&job.ExecutionOptions{ - Priority: 1, - }, "kb.documents.addfile", jobData) - if err != nil { - log.Error("AddFileAsync: Failed to add job execution: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to add job execution: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Step 5: Push the job to execution queue - err = j.Push() - if err != nil { - log.Error("AddFileAsync: Failed to push job: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to push job: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddFileAsync: Job pushed successfully") + log.Info("AddFileAsync: Job created with ID: %s", result.JobID) // Return job_id and doc_id - response.RespondWithSuccess(c, response.StatusCreated, gin.H{ - "job_id": j.JobID, - "doc_id": req.DocID, - }) + response.RespondWithSuccess(c, response.StatusCreated, result) } // ProcessAddFile documents.addfile Knowledge Base add file processor @@ -473,101 +197,141 @@ func ProcessAddFile(process *process.Process) interface{} { // Get parameters reqMap := process.ArgsMap(0) - // Check knowledge base instance - if kb.Instance == nil { - exception.New("knowledge base not initialized", 500).Throw() + // Check knowledge base API + if kb.API == nil { + exception.New("knowledge base API not initialized", 500).Throw() } - // Convert parameters to AddFileRequest structure - req := parseAddFileRequest(reqMap) + // Convert parameters to AddFileParams + params := parseAddFileParams(reqMap) - // Get KB config to check if document exists - config, err := kb.GetConfig() - if err != nil { - exception.New("failed to get KB config: %s", 500, err.Error()).Throw() - } - - // Check if document already exists + // Get context ctx := process.Context if ctx == nil { ctx = context.Background() } - existingDoc, err := config.FindDocument(req.DocID, model.QueryParam{}) - if err != nil || existingDoc == nil { - // Document doesn't exist, create it first (sync scenario) - log.Info("ProcessAddFile: Document %s not found, creating new record", req.DocID) - - // Get job_id from request if provided (for async scenario) - var jobID string - if jid, ok := reqMap["job_id"].(string); ok { - jobID = jid - } - err = CreateDocumentRecord(ctx, authorized.ProcessAuthInfo(process), req, jobID) - if err != nil { - exception.New("failed to create document record: %s", 500, err.Error()).Throw() - } - } else { - log.Info("ProcessAddFile: Document %s already exists, processing content only", req.DocID) - } - - // Process file content - err = HandleFileContent(ctx, req) + // Call kb.API + result, err := kb.API.AddFile(ctx, params) if err != nil { - exception.New("failed to process file: %s", 500, err.Error()).Throw() + exception.New("failed to add file: %s", 500, err.Error()).Throw() } // Return result return maps.MapStrAny{ - "doc_id": req.DocID, + "doc_id": result.DocID, } } -// parseAddFileRequest parses request map into AddFileRequest structure -func parseAddFileRequest(reqMap map[string]interface{}) *AddFileRequest { - req := &AddFileRequest{} +// convertAddFileRequest converts AddFileRequest to kbapi.AddFileParams +func convertAddFileRequest(req *AddFileRequest, authInfo *oauthtypes.AuthorizedInfo) *kbapi.AddFileParams { + params := &kbapi.AddFileParams{ + CollectionID: req.CollectionID, + FileID: req.FileID, + Uploader: req.Uploader, + DocID: req.DocID, + Locale: req.Locale, + Metadata: req.Metadata, + } + + // Convert provider configs + if req.Chunking != nil { + params.Chunking = &kbapi.ProviderConfigParams{ + ProviderID: req.Chunking.ProviderID, + OptionID: req.Chunking.OptionID, + } + } + + if req.Embedding != nil { + params.Embedding = &kbapi.ProviderConfigParams{ + ProviderID: req.Embedding.ProviderID, + OptionID: req.Embedding.OptionID, + } + } + + if req.Extraction != nil { + params.Extraction = &kbapi.ProviderConfigParams{ + ProviderID: req.Extraction.ProviderID, + OptionID: req.Extraction.OptionID, + } + } + + if req.Fetcher != nil { + params.Fetcher = &kbapi.ProviderConfigParams{ + ProviderID: req.Fetcher.ProviderID, + OptionID: req.Fetcher.OptionID, + } + } + + if req.Converter != nil { + params.Converter = &kbapi.ProviderConfigParams{ + ProviderID: req.Converter.ProviderID, + OptionID: req.Converter.OptionID, + } + } + + if req.Job != nil { + params.Job = &kbapi.JobOptionsParams{ + Name: req.Job.Name, + Description: req.Job.Description, + Icon: req.Job.Icon, + Category: req.Job.Category, + } + } + + // Set auth scope + if authInfo != nil { + params.AuthScope = authInfo.WithCreateScope(nil) + } + + return params +} + +// parseAddFileParams parses request map into kbapi.AddFileParams +func parseAddFileParams(reqMap map[string]interface{}) *kbapi.AddFileParams { + params := &kbapi.AddFileParams{} // Required fields if collectionID, ok := reqMap["collection_id"].(string); ok { - req.CollectionID = collectionID + params.CollectionID = collectionID } else { exception.New("collection_id is required", 400).Throw() } if fileID, ok := reqMap["file_id"].(string); ok { - req.FileID = fileID + params.FileID = fileID } else { exception.New("file_id is required", 400).Throw() } // Optional fields if uploader, ok := reqMap["uploader"].(string); ok { - req.Uploader = uploader + params.Uploader = uploader } else { - req.Uploader = "local" // Default to local uploader + params.Uploader = "local" // Default to local uploader } if locale, ok := reqMap["locale"].(string); ok { - req.Locale = locale + params.Locale = locale } if docID, ok := reqMap["doc_id"].(string); ok { - req.DocID = docID + params.DocID = docID } // Generate doc_id if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + if params.DocID == "" { + params.DocID = utils.GenDocIDWithCollectionID(params.CollectionID) } // Handle metadata if metadata, ok := reqMap["metadata"].(map[string]interface{}); ok { - req.Metadata = metadata + params.Metadata = metadata } // Handle chunking configuration if chunkingMap, ok := reqMap["chunking"].(map[string]interface{}); ok { - chunking := &ProviderConfig{} + chunking := &kbapi.ProviderConfigParams{} if providerID, ok := chunkingMap["provider_id"].(string); ok { chunking.ProviderID = providerID } else { @@ -576,14 +340,14 @@ func parseAddFileRequest(reqMap map[string]interface{}) *AddFileRequest { if optionID, ok := chunkingMap["option_id"].(string); ok { chunking.OptionID = optionID } - req.Chunking = chunking + params.Chunking = chunking } else { exception.New("chunking configuration is required", 400).Throw() } // Handle embedding configuration if embeddingMap, ok := reqMap["embedding"].(map[string]interface{}); ok { - embedding := &ProviderConfig{} + embedding := &kbapi.ProviderConfigParams{} if providerID, ok := embeddingMap["provider_id"].(string); ok { embedding.ProviderID = providerID } else { @@ -592,50 +356,50 @@ func parseAddFileRequest(reqMap map[string]interface{}) *AddFileRequest { if optionID, ok := embeddingMap["option_id"].(string); ok { embedding.OptionID = optionID } - req.Embedding = embedding + params.Embedding = embedding } else { exception.New("embedding configuration is required", 400).Throw() } // Handle optional extraction configuration if extractionMap, ok := reqMap["extraction"].(map[string]interface{}); ok { - extraction := &ProviderConfig{} + extraction := &kbapi.ProviderConfigParams{} if providerID, ok := extractionMap["provider_id"].(string); ok { extraction.ProviderID = providerID } if optionID, ok := extractionMap["option_id"].(string); ok { extraction.OptionID = optionID } - req.Extraction = extraction + params.Extraction = extraction } // Handle optional fetcher configuration if fetcherMap, ok := reqMap["fetcher"].(map[string]interface{}); ok { - fetcher := &ProviderConfig{} + fetcher := &kbapi.ProviderConfigParams{} if providerID, ok := fetcherMap["provider_id"].(string); ok { fetcher.ProviderID = providerID } if optionID, ok := fetcherMap["option_id"].(string); ok { fetcher.OptionID = optionID } - req.Fetcher = fetcher + params.Fetcher = fetcher } // Handle optional converter configuration if converterMap, ok := reqMap["converter"].(map[string]interface{}); ok { - converter := &ProviderConfig{} + converter := &kbapi.ProviderConfigParams{} if providerID, ok := converterMap["provider_id"].(string); ok { converter.ProviderID = providerID } if optionID, ok := converterMap["option_id"].(string); ok { converter.OptionID = optionID } - req.Converter = converter + params.Converter = converter } // Handle job options if jobMap, ok := reqMap["job"].(map[string]interface{}); ok { - job := &JobOptions{} + job := &kbapi.JobOptionsParams{} if name, ok := jobMap["name"].(string); ok { job.Name = name } @@ -648,8 +412,35 @@ func parseAddFileRequest(reqMap map[string]interface{}) *AddFileRequest { if category, ok := jobMap["category"].(string); ok { job.Category = category } - req.Job = job + params.Job = job } - return req + return params +} + +// validateFileExists validates that the file exists in the attachment manager +func validateFileExists(c *gin.Context, req *AddFileRequest) error { + // Get file manager + m, ok := attachment.Managers[req.Uploader] + if !ok { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Invalid uploader: " + req.Uploader + " not found", + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return fmt.Errorf("invalid uploader: %s not found", req.Uploader) + } + + // Check if the file exists + exists := m.Exists(c.Request.Context(), req.FileID) + if !exists { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "File not found: " + req.FileID, + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return fmt.Errorf("file not found: %s", req.FileID) + } + + return nil } diff --git a/openapi/kb/addtext.go b/openapi/kb/addtext.go index 0a483ed4..e34b7674 100644 --- a/openapi/kb/addtext.go +++ b/openapi/kb/addtext.go @@ -2,179 +2,26 @@ package kb import ( "context" - "fmt" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/utils" - "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/job" "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" ) -// CreateTextDocumentRecord creates a text document record in the database immediately -// This is called synchronously when the API request comes in -func CreateTextDocumentRecord(ctx context.Context, req *AddTextRequest, jobID string) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Prepare document data for database - documentData := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": "Text Document", - "type": "text", - "status": "pending", - "text_content": req.Text, - "size": int64(len(req.Text)), - "job_id": jobID, - } - - // Use title from metadata if available - if req.Metadata != nil { - if title, ok := req.Metadata["title"].(string); ok && title != "" { - documentData["name"] = title - } - } - - // Add base request fields - req.BaseUpsertRequest.AddBaseFields(documentData) - - // Create database record - _, err = config.CreateDocument(maps.MapStrAny(documentData)) - if err != nil { - return fmt.Errorf("failed to save document metadata: %w", err) - } - - return nil -} - -// HandleTextContent processes the actual text content and updates the knowledge base -// This is called asynchronously by the job system -func HandleTextContent(ctx context.Context, req *AddTextRequest) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Convert request to UpsertOptions - upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions() - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to convert request to upsert options: %w", err) - } - - // Perform upsert operation with text - _, err = kb.Instance.AddText(ctx, req.Text, upsertOptions) - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to add text: %w", err) - } - - // Update status to completed after successful processing - if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { - log.Error("Failed to update document status to completed: %v", err) - } - - // Update segment count for the document - if segmentCount, err := kb.Instance.SegmentCount(ctx, req.DocID); err != nil { - log.Error("Failed to get segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Got segment count %d for document %s", segmentCount, req.DocID) - if err := config.UpdateSegmentCount(req.DocID, segmentCount); err != nil { - log.Error("Failed to update segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Successfully updated segment count to %d for document %s", segmentCount, req.DocID) - } - } - - // Update document count for the collection and sync to GraphRag - if err := UpdateDocumentCountWithSync(req.CollectionID, config); err != nil { - log.Error("Failed to update document count for collection %s: %v", req.CollectionID, err) - } else { - log.Info("Successfully updated document count for collection %s", req.CollectionID) - } - - return nil -} - -// AddTextHandler processes a text addition request with business logic only -// This function combines both document creation and content processing for sync operations -func AddTextHandler(ctx context.Context, req *AddTextRequest, jobID ...string) error { - // Validate request - if err := req.Validate(); err != nil { - return err - } - - // DocID should be generated by the caller before calling this function - if req.DocID == "" { - return fmt.Errorf("document ID is required") - } - - // For sync operations, create document record and process content immediately - var jid string - if len(jobID) > 0 { - jid = jobID[0] - } - - // Create document record - if err := CreateTextDocumentRecord(ctx, req, jid); err != nil { - return err - } - - // Process text content - return HandleTextContent(ctx, req) -} - -// addTextWithRequest processes a text addition with pre-parsed request using Gin context -func addTextWithRequest(c *gin.Context, req *AddTextRequest) { - // Use the business logic function - err := AddTextHandler(c.Request.Context(), req) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Return success response - result := gin.H{ - "message": "Text added successfully", - "collection_id": req.CollectionID, - "doc_id": req.DocID, - } - - response.RespondWithSuccess(c, response.StatusCreated, result) -} - -// AddText adds text to a collection +// AddText adds text to a collection (sync) func AddText(c *gin.Context) { var req AddTextRequest - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -203,8 +50,44 @@ func AddText(c *gin.Context) { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } - // Process the request - addTextWithRequest(c, &req) + // Check collection permission + authInfo := authorized.GetInfo(c) + hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // 403 Forbidden + if !hasPermission { + errorResp := &response.ErrorResponse{ + Code: response.ErrAccessDenied.Code, + ErrorDescription: "Forbidden: No permission to update collection", + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // Convert request to API params + params := convertAddTextRequest(&req, authInfo) + + // Call kb.API + result, err := kb.API.AddText(c.Request.Context(), params) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Return success response + response.RespondWithSuccess(c, response.StatusCreated, result) } // AddTextAsync adds text to a collection asynchronously @@ -213,9 +96,9 @@ func AddTextAsync(c *gin.Context) { log.Info("AddTextAsync: Starting async text addition") - // Check if kb.Instance is available - if !checkKBInstance(c) { - log.Error("AddTextAsync: KB instance check failed") + // Check if kb.API is available + if !checkKBAPI(c) { + log.Error("AddTextAsync: KB API check failed") return } @@ -230,7 +113,7 @@ func AddTextAsync(c *gin.Context) { return } - log.Info("AddTextAsync: Request parsed successfully: %+v", req) + log.Info("AddTextAsync: Request parsed successfully") // Validate request if err := req.Validate(); err != nil { @@ -245,15 +128,6 @@ func AddTextAsync(c *gin.Context) { log.Info("AddTextAsync: Request validation passed") - // Convert request to UpsertOptions (just for validation) - _, err := getUpsertOptions(c, &req.BaseUpsertRequest) - if err != nil { - log.Error("AddTextAsync: UpsertOptions validation failed: %v", err) - return - } - - log.Info("AddTextAsync: UpsertOptions validation passed") - // Generate document ID if not provided if req.DocID == "" { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) @@ -261,109 +135,50 @@ func AddTextAsync(c *gin.Context) { log.Info("AddTextAsync: Generated doc_id: %s", req.DocID) - // Step 1: Get job options with defaults - jobName, jobDescription, jobIcon, jobCategory := req.GetJobOptions( - "Knowledge Base Text Processing", // default name - "Processing and indexing text content for knowledge base search", // default description - "library_add", // default icon (Material Icon) - "Knowledge Base", // default category - ) - - // Create job data - jobCreateData := map[string]interface{}{ - "name": jobName, - "description": jobDescription, - "category_name": jobCategory, // Pass category name directly, let SaveJob handle it - } - if jobIcon != "" { - jobCreateData["icon"] = jobIcon - } - - // Create and save Job in one step to get JobID - j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + // Check collection permission + authInfo := authorized.GetInfo(c) + hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) if err != nil { - log.Error("AddTextAsync: Job creation and save failed: %v", err) errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create and save job: " + err.Error(), + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // 403 Forbidden + if !hasPermission { + errorResp := &response.ErrorResponse{ + Code: response.ErrAccessDenied.Code, + ErrorDescription: "Forbidden: No permission to update collection", + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // Convert request to API params + params := convertAddTextRequest(&req, authInfo) + + // Call kb.API async + result, err := kb.API.AddTextAsync(c.Request.Context(), params) + if err != nil { + log.Error("AddTextAsync: Failed to add text async: %v", err) + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } - log.Info("AddTextAsync: Job created and saved with ID: %s", j.JobID) - - // Step 2: Create document record immediately with job_id - err = CreateTextDocumentRecord(c.Request.Context(), &req, j.JobID) - if err != nil { - log.Error("AddTextAsync: Failed to create document record: %v", err) - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create document record: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddTextAsync: Document record created successfully") - - // Step 3: Add execution to job - jobData := map[string]interface{}{ - "collection_id": req.CollectionID, - "text": req.Text, - "locale": req.Locale, - "doc_id": req.DocID, - "metadata": req.Metadata, - "chunking": req.Chunking, - "embedding": req.Embedding, - "extraction": req.Extraction, - "fetcher": req.Fetcher, - "converter": req.Converter, - } - - err = j.Add(&job.ExecutionOptions{ - Priority: 1, - }, "kb.documents.addtext", jobData) - if err != nil { - log.Error("AddTextAsync: Failed to add job execution: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to add job execution: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Step 4: Push the job to execution queue - err = j.Push() - if err != nil { - log.Error("AddTextAsync: Failed to push job: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to push job: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddTextAsync: Job pushed successfully") + log.Info("AddTextAsync: Job created with ID: %s", result.JobID) // Return job_id and doc_id - response.RespondWithSuccess(c, response.StatusCreated, gin.H{ - "job_id": j.JobID, - "doc_id": req.DocID, - }) + response.RespondWithSuccess(c, response.StatusCreated, result) } -// ProcessAddText documents.addtext Knowledge Base add text processor (sync version) +// ProcessAddText documents.addtext Knowledge Base add text processor // Args[0] map: Request parameters {"collection_id": "collection", "text": "content", ...} // Return: map: Response data {"doc_id": "document_id"} func ProcessAddText(process *process.Process) interface{} { @@ -372,96 +187,134 @@ func ProcessAddText(process *process.Process) interface{} { // Get parameters reqMap := process.ArgsMap(0) - // Check knowledge base instance - if kb.Instance == nil { - exception.New("knowledge base not initialized", 500).Throw() + // Check knowledge base API + if kb.API == nil { + exception.New("knowledge base API not initialized", 500).Throw() } - // Convert parameters to AddTextRequest structure - req := parseAddTextRequest(reqMap) + // Convert parameters to AddTextParams + params := parseAddTextParams(reqMap) - // Get KB config to check if document exists - config, err := kb.GetConfig() - if err != nil { - exception.New("failed to get KB config: %s", 500, err.Error()).Throw() - } - - // Check if document already exists + // Get context ctx := process.Context if ctx == nil { ctx = context.Background() } - existingDoc, err := config.FindDocument(req.DocID, model.QueryParam{}) - if err != nil || existingDoc == nil { - // Document doesn't exist, create it first (sync scenario) - log.Info("ProcessAddText: Document %s not found, creating new record", req.DocID) - - // Get job_id from request if provided (for async scenario) - var jobID string - if jid, ok := reqMap["job_id"].(string); ok { - jobID = jid - } - - err = CreateTextDocumentRecord(ctx, req, jobID) - if err != nil { - exception.New("failed to create document record: %s", 500, err.Error()).Throw() - } - } else { - log.Info("ProcessAddText: Document %s already exists, processing content only", req.DocID) - } - - // Process text content - err = HandleTextContent(ctx, req) + // Call kb.API + result, err := kb.API.AddText(ctx, params) if err != nil { - exception.New("failed to process text: %s", 500, err.Error()).Throw() + exception.New("failed to add text: %s", 500, err.Error()).Throw() } // Return result return maps.MapStrAny{ - "doc_id": req.DocID, + "doc_id": result.DocID, } } -// parseAddTextRequest parses request map into AddTextRequest structure -func parseAddTextRequest(reqMap map[string]interface{}) *AddTextRequest { - req := &AddTextRequest{} +// convertAddTextRequest converts AddTextRequest to kbapi.AddTextParams +func convertAddTextRequest(req *AddTextRequest, authInfo *oauthtypes.AuthorizedInfo) *kbapi.AddTextParams { + params := &kbapi.AddTextParams{ + CollectionID: req.CollectionID, + Text: req.Text, + DocID: req.DocID, + Locale: req.Locale, + Metadata: req.Metadata, + } + + // Convert provider configs + if req.Chunking != nil { + params.Chunking = &kbapi.ProviderConfigParams{ + ProviderID: req.Chunking.ProviderID, + OptionID: req.Chunking.OptionID, + } + } + + if req.Embedding != nil { + params.Embedding = &kbapi.ProviderConfigParams{ + ProviderID: req.Embedding.ProviderID, + OptionID: req.Embedding.OptionID, + } + } + + if req.Extraction != nil { + params.Extraction = &kbapi.ProviderConfigParams{ + ProviderID: req.Extraction.ProviderID, + OptionID: req.Extraction.OptionID, + } + } + + if req.Fetcher != nil { + params.Fetcher = &kbapi.ProviderConfigParams{ + ProviderID: req.Fetcher.ProviderID, + OptionID: req.Fetcher.OptionID, + } + } + + if req.Converter != nil { + params.Converter = &kbapi.ProviderConfigParams{ + ProviderID: req.Converter.ProviderID, + OptionID: req.Converter.OptionID, + } + } + + if req.Job != nil { + params.Job = &kbapi.JobOptionsParams{ + Name: req.Job.Name, + Description: req.Job.Description, + Icon: req.Job.Icon, + Category: req.Job.Category, + } + } + + // Set auth scope + if authInfo != nil { + params.AuthScope = authInfo.WithCreateScope(nil) + } + + return params +} + +// parseAddTextParams parses request map into kbapi.AddTextParams +func parseAddTextParams(reqMap map[string]interface{}) *kbapi.AddTextParams { + params := &kbapi.AddTextParams{} // Required fields if collectionID, ok := reqMap["collection_id"].(string); ok { - req.CollectionID = collectionID + params.CollectionID = collectionID } else { exception.New("collection_id is required", 400).Throw() } if text, ok := reqMap["text"].(string); ok { - req.Text = text + params.Text = text } else { exception.New("text is required", 400).Throw() } // Optional fields if locale, ok := reqMap["locale"].(string); ok { - req.Locale = locale + params.Locale = locale } if docID, ok := reqMap["doc_id"].(string); ok { - req.DocID = docID + params.DocID = docID } // Generate doc_id if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + if params.DocID == "" { + params.DocID = utils.GenDocIDWithCollectionID(params.CollectionID) } // Handle metadata if metadata, ok := reqMap["metadata"].(map[string]interface{}); ok { - req.Metadata = metadata + params.Metadata = metadata } // Handle chunking configuration if chunkingMap, ok := reqMap["chunking"].(map[string]interface{}); ok { - chunking := &ProviderConfig{} + chunking := &kbapi.ProviderConfigParams{} if providerID, ok := chunkingMap["provider_id"].(string); ok { chunking.ProviderID = providerID } else { @@ -470,14 +323,14 @@ func parseAddTextRequest(reqMap map[string]interface{}) *AddTextRequest { if optionID, ok := chunkingMap["option_id"].(string); ok { chunking.OptionID = optionID } - req.Chunking = chunking + params.Chunking = chunking } else { exception.New("chunking configuration is required", 400).Throw() } // Handle embedding configuration if embeddingMap, ok := reqMap["embedding"].(map[string]interface{}); ok { - embedding := &ProviderConfig{} + embedding := &kbapi.ProviderConfigParams{} if providerID, ok := embeddingMap["provider_id"].(string); ok { embedding.ProviderID = providerID } else { @@ -486,50 +339,50 @@ func parseAddTextRequest(reqMap map[string]interface{}) *AddTextRequest { if optionID, ok := embeddingMap["option_id"].(string); ok { embedding.OptionID = optionID } - req.Embedding = embedding + params.Embedding = embedding } else { exception.New("embedding configuration is required", 400).Throw() } // Handle optional extraction configuration if extractionMap, ok := reqMap["extraction"].(map[string]interface{}); ok { - extraction := &ProviderConfig{} + extraction := &kbapi.ProviderConfigParams{} if providerID, ok := extractionMap["provider_id"].(string); ok { extraction.ProviderID = providerID } if optionID, ok := extractionMap["option_id"].(string); ok { extraction.OptionID = optionID } - req.Extraction = extraction + params.Extraction = extraction } // Handle optional fetcher configuration if fetcherMap, ok := reqMap["fetcher"].(map[string]interface{}); ok { - fetcher := &ProviderConfig{} + fetcher := &kbapi.ProviderConfigParams{} if providerID, ok := fetcherMap["provider_id"].(string); ok { fetcher.ProviderID = providerID } if optionID, ok := fetcherMap["option_id"].(string); ok { fetcher.OptionID = optionID } - req.Fetcher = fetcher + params.Fetcher = fetcher } // Handle optional converter configuration if converterMap, ok := reqMap["converter"].(map[string]interface{}); ok { - converter := &ProviderConfig{} + converter := &kbapi.ProviderConfigParams{} if providerID, ok := converterMap["provider_id"].(string); ok { converter.ProviderID = providerID } if optionID, ok := converterMap["option_id"].(string); ok { converter.OptionID = optionID } - req.Converter = converter + params.Converter = converter } // Handle job options if jobMap, ok := reqMap["job"].(map[string]interface{}); ok { - job := &JobOptions{} + job := &kbapi.JobOptionsParams{} if name, ok := jobMap["name"].(string); ok { job.Name = name } @@ -542,8 +395,8 @@ func parseAddTextRequest(reqMap map[string]interface{}) *AddTextRequest { if category, ok := jobMap["category"].(string); ok { job.Category = category } - req.Job = job + params.Job = job } - return req + return params } diff --git a/openapi/kb/addurl.go b/openapi/kb/addurl.go index 5dff7651..02686b8c 100644 --- a/openapi/kb/addurl.go +++ b/openapi/kb/addurl.go @@ -2,179 +2,26 @@ package kb import ( "context" - "fmt" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/utils" - "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/job" "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" ) -// CreateURLDocumentRecord creates a URL document record in the database immediately -// This is called synchronously when the API request comes in -func CreateURLDocumentRecord(ctx context.Context, req *AddURLRequest, jobID string) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Prepare document data for database - documentData := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": "URL Document", - "type": "url", - "status": "pending", - "url": req.URL, - "job_id": jobID, - } - - // Use title from metadata if available - if req.Metadata != nil { - if title, ok := req.Metadata["title"].(string); ok && title != "" { - documentData["name"] = title - } - } - - // Add base request fields - req.BaseUpsertRequest.AddBaseFields(documentData) - - // Create database record - _, err = config.CreateDocument(maps.MapStrAny(documentData)) - if err != nil { - return fmt.Errorf("failed to save document metadata: %w", err) - } - - return nil -} - -// HandleURLContent processes the actual URL content and updates the knowledge base -// This is called asynchronously by the job system -func HandleURLContent(ctx context.Context, req *AddURLRequest) error { - // Check if kb.Instance is available - if kb.Instance == nil { - return fmt.Errorf("knowledge base not initialized") - } - - // Get KB config - config, err := kb.GetConfig() - if err != nil { - return fmt.Errorf("failed to get KB config: %w", err) - } - - // Convert request to UpsertOptions - upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions() - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to convert request to upsert options: %w", err) - } - - // Perform upsert operation with URL - _, err = kb.Instance.AddURL(ctx, req.URL, upsertOptions) - if err != nil { - // Update status to error - config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - return fmt.Errorf("failed to add URL: %w", err) - } - - // Update status to completed after successful processing - if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { - log.Error("Failed to update document status to completed: %v", err) - } - - // Update segment count for the document - if segmentCount, err := kb.Instance.SegmentCount(ctx, req.DocID); err != nil { - log.Error("Failed to get segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Got segment count %d for document %s", segmentCount, req.DocID) - if err := config.UpdateSegmentCount(req.DocID, segmentCount); err != nil { - log.Error("Failed to update segment count for document %s: %v", req.DocID, err) - } else { - log.Info("Successfully updated segment count to %d for document %s", segmentCount, req.DocID) - } - } - - // Update document count for the collection and sync to GraphRag - if err := UpdateDocumentCountWithSync(req.CollectionID, config); err != nil { - log.Error("Failed to update document count for collection %s: %v", req.CollectionID, err) - } else { - log.Info("Successfully updated document count for collection %s", req.CollectionID) - } - - return nil -} - -// AddURLHandler processes a URL addition request with business logic only -// This function combines both document creation and content processing for sync operations -func AddURLHandler(ctx context.Context, req *AddURLRequest, jobID ...string) error { - // Validate request - if err := req.Validate(); err != nil { - return err - } - - // DocID should be generated by the caller before calling this function - if req.DocID == "" { - return fmt.Errorf("document ID is required") - } - - // For sync operations, create document record and process content immediately - var jid string - if len(jobID) > 0 { - jid = jobID[0] - } - - // Create document record - if err := CreateURLDocumentRecord(ctx, req, jid); err != nil { - return err - } - - // Process URL content - return HandleURLContent(ctx, req) -} - -// addURLWithRequest processes a URL addition with pre-parsed request using Gin context -func addURLWithRequest(c *gin.Context, req *AddURLRequest) { - // Use the business logic function - err := AddURLHandler(c.Request.Context(), req) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Return success response - result := gin.H{ - "message": "URL added successfully", - "collection_id": req.CollectionID, - "url": req.URL, - "doc_id": req.DocID, - } - - response.RespondWithSuccess(c, response.StatusCreated, result) -} - -// AddURL adds a URL to a collection +// AddURL adds a URL to a collection (sync) func AddURL(c *gin.Context) { var req AddURLRequest - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -203,8 +50,44 @@ func AddURL(c *gin.Context) { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } - // Process the request - addURLWithRequest(c, &req) + // Check collection permission + authInfo := authorized.GetInfo(c) + hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // 403 Forbidden + if !hasPermission { + errorResp := &response.ErrorResponse{ + Code: response.ErrAccessDenied.Code, + ErrorDescription: "Forbidden: No permission to update collection", + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // Convert request to API params + params := convertAddURLRequest(&req, authInfo) + + // Call kb.API + result, err := kb.API.AddURL(c.Request.Context(), params) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Return success response + response.RespondWithSuccess(c, response.StatusCreated, result) } // AddURLAsync adds a URL to a collection asynchronously @@ -213,9 +96,9 @@ func AddURLAsync(c *gin.Context) { log.Info("AddURLAsync: Starting async URL addition") - // Check if kb.Instance is available - if !checkKBInstance(c) { - log.Error("AddURLAsync: KB instance check failed") + // Check if kb.API is available + if !checkKBAPI(c) { + log.Error("AddURLAsync: KB API check failed") return } @@ -230,7 +113,7 @@ func AddURLAsync(c *gin.Context) { return } - log.Info("AddURLAsync: Request parsed successfully: %+v", req) + log.Info("AddURLAsync: Request parsed successfully") // Validate request if err := req.Validate(); err != nil { @@ -245,15 +128,6 @@ func AddURLAsync(c *gin.Context) { log.Info("AddURLAsync: Request validation passed") - // Convert request to UpsertOptions (just for validation) - _, err := getUpsertOptions(c, &req.BaseUpsertRequest) - if err != nil { - log.Error("AddURLAsync: UpsertOptions validation failed: %v", err) - return - } - - log.Info("AddURLAsync: UpsertOptions validation passed") - // Generate document ID if not provided if req.DocID == "" { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) @@ -261,109 +135,50 @@ func AddURLAsync(c *gin.Context) { log.Info("AddURLAsync: Generated doc_id: %s", req.DocID) - // Step 1: Get job options with defaults - jobName, jobDescription, jobIcon, jobCategory := req.GetJobOptions( - "Knowledge Base Web Content Processing", // default name - "Fetching and indexing web content for knowledge base search", // default description - "library_add", // default icon (Material Icon) - "Knowledge Base", // default category - ) - - // Create job data - jobCreateData := map[string]interface{}{ - "name": jobName, - "description": jobDescription, - "category_name": jobCategory, // Pass category name directly, let SaveJob handle it - } - if jobIcon != "" { - jobCreateData["icon"] = jobIcon - } - - // Create and save Job in one step to get JobID - j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData) + // Check collection permission + authInfo := authorized.GetInfo(c) + hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID) if err != nil { - log.Error("AddURLAsync: Job creation and save failed: %v", err) errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create and save job: " + err.Error(), + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // 403 Forbidden + if !hasPermission { + errorResp := &response.ErrorResponse{ + Code: response.ErrAccessDenied.Code, + ErrorDescription: "Forbidden: No permission to update collection", + } + response.RespondWithError(c, response.StatusForbidden, errorResp) + return + } + + // Convert request to API params + params := convertAddURLRequest(&req, authInfo) + + // Call kb.API async + result, err := kb.API.AddURLAsync(c.Request.Context(), params) + if err != nil { + log.Error("AddURLAsync: Failed to add URL async: %v", err) + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } - log.Info("AddURLAsync: Job created and saved with ID: %s", j.JobID) - - // Step 2: Create document record immediately with job_id - err = CreateURLDocumentRecord(c.Request.Context(), &req, j.JobID) - if err != nil { - log.Error("AddURLAsync: Failed to create document record: %v", err) - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to create document record: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddURLAsync: Document record created successfully") - - // Step 3: Add execution to job - jobData := map[string]interface{}{ - "collection_id": req.CollectionID, - "url": req.URL, - "locale": req.Locale, - "doc_id": req.DocID, - "metadata": req.Metadata, - "chunking": req.Chunking, - "embedding": req.Embedding, - "extraction": req.Extraction, - "fetcher": req.Fetcher, - "converter": req.Converter, - } - - err = j.Add(&job.ExecutionOptions{ - Priority: 1, - }, "kb.documents.addurl", jobData) - if err != nil { - log.Error("AddURLAsync: Failed to add job execution: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to add job execution: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - // Step 4: Push the job to execution queue - err = j.Push() - if err != nil { - log.Error("AddURLAsync: Failed to push job: %v", err) - // Rollback: remove document record - if config, err := kb.GetConfig(); err == nil { - config.RemoveDocument(req.DocID) - } - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to push job: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - - log.Info("AddURLAsync: Job pushed successfully") + log.Info("AddURLAsync: Job created with ID: %s", result.JobID) // Return job_id and doc_id - response.RespondWithSuccess(c, response.StatusCreated, gin.H{ - "job_id": j.JobID, - "doc_id": req.DocID, - }) + response.RespondWithSuccess(c, response.StatusCreated, result) } -// ProcessAddURL documents.addurl Knowledge Base add URL processor (sync version) +// ProcessAddURL documents.addurl Knowledge Base add URL processor // Args[0] map: Request parameters {"collection_id": "collection", "url": "https://example.com", ...} // Return: map: Response data {"doc_id": "document_id"} func ProcessAddURL(process *process.Process) interface{} { @@ -372,96 +187,134 @@ func ProcessAddURL(process *process.Process) interface{} { // Get parameters reqMap := process.ArgsMap(0) - // Check knowledge base instance - if kb.Instance == nil { - exception.New("knowledge base not initialized", 500).Throw() + // Check knowledge base API + if kb.API == nil { + exception.New("knowledge base API not initialized", 500).Throw() } - // Convert parameters to AddURLRequest structure - req := parseAddURLRequest(reqMap) + // Convert parameters to AddURLParams + params := parseAddURLParams(reqMap) - // Get KB config to check if document exists - config, err := kb.GetConfig() - if err != nil { - exception.New("failed to get KB config: %s", 500, err.Error()).Throw() - } - - // Check if document already exists + // Get context ctx := process.Context if ctx == nil { ctx = context.Background() } - existingDoc, err := config.FindDocument(req.DocID, model.QueryParam{}) - if err != nil || existingDoc == nil { - // Document doesn't exist, create it first (sync scenario) - log.Info("ProcessAddURL: Document %s not found, creating new record", req.DocID) - - // Get job_id from request if provided (for async scenario) - var jobID string - if jid, ok := reqMap["job_id"].(string); ok { - jobID = jid - } - - err = CreateURLDocumentRecord(ctx, req, jobID) - if err != nil { - exception.New("failed to create document record: %s", 500, err.Error()).Throw() - } - } else { - log.Info("ProcessAddURL: Document %s already exists, processing content only", req.DocID) - } - - // Process URL content - err = HandleURLContent(ctx, req) + // Call kb.API + result, err := kb.API.AddURL(ctx, params) if err != nil { - exception.New("failed to process URL: %s", 500, err.Error()).Throw() + exception.New("failed to add URL: %s", 500, err.Error()).Throw() } // Return result return maps.MapStrAny{ - "doc_id": req.DocID, + "doc_id": result.DocID, } } -// parseAddURLRequest parses request map into AddURLRequest structure -func parseAddURLRequest(reqMap map[string]interface{}) *AddURLRequest { - req := &AddURLRequest{} +// convertAddURLRequest converts AddURLRequest to kbapi.AddURLParams +func convertAddURLRequest(req *AddURLRequest, authInfo *oauthtypes.AuthorizedInfo) *kbapi.AddURLParams { + params := &kbapi.AddURLParams{ + CollectionID: req.CollectionID, + URL: req.URL, + DocID: req.DocID, + Locale: req.Locale, + Metadata: req.Metadata, + } + + // Convert provider configs + if req.Chunking != nil { + params.Chunking = &kbapi.ProviderConfigParams{ + ProviderID: req.Chunking.ProviderID, + OptionID: req.Chunking.OptionID, + } + } + + if req.Embedding != nil { + params.Embedding = &kbapi.ProviderConfigParams{ + ProviderID: req.Embedding.ProviderID, + OptionID: req.Embedding.OptionID, + } + } + + if req.Extraction != nil { + params.Extraction = &kbapi.ProviderConfigParams{ + ProviderID: req.Extraction.ProviderID, + OptionID: req.Extraction.OptionID, + } + } + + if req.Fetcher != nil { + params.Fetcher = &kbapi.ProviderConfigParams{ + ProviderID: req.Fetcher.ProviderID, + OptionID: req.Fetcher.OptionID, + } + } + + if req.Converter != nil { + params.Converter = &kbapi.ProviderConfigParams{ + ProviderID: req.Converter.ProviderID, + OptionID: req.Converter.OptionID, + } + } + + if req.Job != nil { + params.Job = &kbapi.JobOptionsParams{ + Name: req.Job.Name, + Description: req.Job.Description, + Icon: req.Job.Icon, + Category: req.Job.Category, + } + } + + // Set auth scope + if authInfo != nil { + params.AuthScope = authInfo.WithCreateScope(nil) + } + + return params +} + +// parseAddURLParams parses request map into kbapi.AddURLParams +func parseAddURLParams(reqMap map[string]interface{}) *kbapi.AddURLParams { + params := &kbapi.AddURLParams{} // Required fields if collectionID, ok := reqMap["collection_id"].(string); ok { - req.CollectionID = collectionID + params.CollectionID = collectionID } else { exception.New("collection_id is required", 400).Throw() } if url, ok := reqMap["url"].(string); ok { - req.URL = url + params.URL = url } else { exception.New("url is required", 400).Throw() } // Optional fields if locale, ok := reqMap["locale"].(string); ok { - req.Locale = locale + params.Locale = locale } if docID, ok := reqMap["doc_id"].(string); ok { - req.DocID = docID + params.DocID = docID } // Generate doc_id if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + if params.DocID == "" { + params.DocID = utils.GenDocIDWithCollectionID(params.CollectionID) } // Handle metadata if metadata, ok := reqMap["metadata"].(map[string]interface{}); ok { - req.Metadata = metadata + params.Metadata = metadata } // Handle chunking configuration if chunkingMap, ok := reqMap["chunking"].(map[string]interface{}); ok { - chunking := &ProviderConfig{} + chunking := &kbapi.ProviderConfigParams{} if providerID, ok := chunkingMap["provider_id"].(string); ok { chunking.ProviderID = providerID } else { @@ -470,14 +323,14 @@ func parseAddURLRequest(reqMap map[string]interface{}) *AddURLRequest { if optionID, ok := chunkingMap["option_id"].(string); ok { chunking.OptionID = optionID } - req.Chunking = chunking + params.Chunking = chunking } else { exception.New("chunking configuration is required", 400).Throw() } // Handle embedding configuration if embeddingMap, ok := reqMap["embedding"].(map[string]interface{}); ok { - embedding := &ProviderConfig{} + embedding := &kbapi.ProviderConfigParams{} if providerID, ok := embeddingMap["provider_id"].(string); ok { embedding.ProviderID = providerID } else { @@ -486,50 +339,50 @@ func parseAddURLRequest(reqMap map[string]interface{}) *AddURLRequest { if optionID, ok := embeddingMap["option_id"].(string); ok { embedding.OptionID = optionID } - req.Embedding = embedding + params.Embedding = embedding } else { exception.New("embedding configuration is required", 400).Throw() } // Handle optional extraction configuration if extractionMap, ok := reqMap["extraction"].(map[string]interface{}); ok { - extraction := &ProviderConfig{} + extraction := &kbapi.ProviderConfigParams{} if providerID, ok := extractionMap["provider_id"].(string); ok { extraction.ProviderID = providerID } if optionID, ok := extractionMap["option_id"].(string); ok { extraction.OptionID = optionID } - req.Extraction = extraction + params.Extraction = extraction } // Handle optional fetcher configuration if fetcherMap, ok := reqMap["fetcher"].(map[string]interface{}); ok { - fetcher := &ProviderConfig{} + fetcher := &kbapi.ProviderConfigParams{} if providerID, ok := fetcherMap["provider_id"].(string); ok { fetcher.ProviderID = providerID } if optionID, ok := fetcherMap["option_id"].(string); ok { fetcher.OptionID = optionID } - req.Fetcher = fetcher + params.Fetcher = fetcher } // Handle optional converter configuration if converterMap, ok := reqMap["converter"].(map[string]interface{}); ok { - converter := &ProviderConfig{} + converter := &kbapi.ProviderConfigParams{} if providerID, ok := converterMap["provider_id"].(string); ok { converter.ProviderID = providerID } if optionID, ok := converterMap["option_id"].(string); ok { converter.OptionID = optionID } - req.Converter = converter + params.Converter = converter } // Handle job options if jobMap, ok := reqMap["job"].(map[string]interface{}); ok { - job := &JobOptions{} + job := &kbapi.JobOptionsParams{} if name, ok := jobMap["name"].(string); ok { job.Name = name } @@ -542,8 +395,8 @@ func parseAddURLRequest(reqMap map[string]interface{}) *AddURLRequest { if category, ok := jobMap["category"].(string); ok { job.Category = category } - req.Job = job + params.Job = job } - return req + return params } diff --git a/openapi/kb/document.go b/openapi/kb/document.go index 407dde02..826a5832 100644 --- a/openapi/kb/document.go +++ b/openapi/kb/document.go @@ -6,59 +6,19 @@ import ( "strings" "github.com/gin-gonic/gin" - "github.com/yaoapp/gou/graphrag/types" - kbutils "github.com/yaoapp/gou/graphrag/utils" "github.com/yaoapp/gou/model" - "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" + kbapi "github.com/yaoapp/yao/kb/api" "github.com/yaoapp/yao/openapi/oauth/authorized" "github.com/yaoapp/yao/openapi/response" ) -// Document field definitions -var ( - // availableDocumentFields defines all available fields for security filtering - availableDocumentFields = map[string]bool{ - "id": true, "document_id": true, "collection_id": true, "name": true, - "description": true, "status": true, "type": true, "size": true, - "segment_count": true, "job_id": true, "uploader_id": true, "tags": true, - "locale": true, "system": true, "readonly": true, "sort": true, "cover": true, - "file_id": true, "file_name": true, "file_mime_type": true, - "url": true, "url_title": true, "text_content": true, - "converter_provider_id": true, "converter_option_id": true, "converter_properties": true, - "fetcher_provider_id": true, "fetcher_option_id": true, "fetcher_properties": true, - "chunking_provider_id": true, "chunking_option_id": true, "chunking_properties": true, - "extraction_provider_id": true, "extraction_option_id": true, "extraction_properties": true, - "processed_at": true, "error_message": true, "created_at": true, "updated_at": true, - } - - // defaultDocumentFields defines the default compact field list - defaultDocumentFields = []interface{}{ - "id", "document_id", "collection_id", "name", "description", - "cover", "tags", "type", "size", "segment_count", "status", "locale", - "system", "readonly", "file_id", "file_name", "file_mime_type", "uploader_id", - "url", "url_title", "text_content", // 添加 URL 和文本内容字段 - "error_message", "created_at", "updated_at", - } - - // validSortFields defines valid fields for sorting - validSortFields = map[string]bool{ - "created_at": true, - "updated_at": true, - "name": true, - "size": true, - "segment_count": true, - "sort": true, - "processed_at": true, - } -) - // Document Management Handlers // ListDocuments lists documents with pagination func ListDocuments(c *gin.Context) { - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -77,70 +37,39 @@ func ListDocuments(c *gin.Context) { } } - // Get KB instance and config - kbInstance := kb.Instance.(*kb.KnowledgeBase) - config := kbInstance.Config - // 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 != "" && availableDocumentFields[field] { + if field != "" && kbapi.AvailableDocumentFields[field] { selectFields = append(selectFields, field) } } // If no valid fields found, use default if len(selectFields) == 0 { - selectFields = defaultDocumentFields + selectFields = kbapi.DefaultDocumentFields } } else { - selectFields = defaultDocumentFields - } - - // Build query parameters - param := model.QueryParam{ - Select: selectFields, - } - - // Add filters - var wheres []model.QueryWhere - - // Filter by keywords (search in name and description) - if keywords := strings.TrimSpace(c.Query("keywords")); keywords != "" { - wheres = append(wheres, model.QueryWhere{ - Column: "name", - Value: "%" + keywords + "%", - OP: "like", - }) - wheres = append(wheres, model.QueryWhere{ - Column: "description", - Value: "%" + keywords + "%", - OP: "like", - Wheres: []model.QueryWhere{}, - Method: "orwhere", - }) - } - - // Filter by tag - if tag := strings.TrimSpace(c.Query("tag")); tag != "" { - wheres = append(wheres, model.QueryWhere{ - Column: "tags", - Value: "%" + tag + "%", - OP: "like", - }) + selectFields = kbapi.DefaultDocumentFields } // Get authorized information authInfo := authorized.GetInfo(c) + // Build filter for kb.API + filter := &kbapi.ListDocumentsFilter{ + Page: page, + PageSize: pagesize, + Keywords: strings.TrimSpace(c.Query("keywords")), + Tag: strings.TrimSpace(c.Query("tag")), + Select: selectFields, + } + // Filter by collection_id - // If collection_id is provided, validate collection permission - // If not provided, filter by authorization constraints (TeamOnly or OwnerOnly) collectionID := strings.TrimSpace(c.Query("collection_id")) if collectionID != "" { - // Validate collection permission hasPermission, err := checkCollectionPermission(authInfo, collectionID, true) if err != nil { @@ -156,84 +85,44 @@ func ListDocuments(c *gin.Context) { if !hasPermission { errorResp := &response.ErrorResponse{ Code: response.ErrAccessDenied.Code, - ErrorDescription: "Forbidden: No permission to update collection", + ErrorDescription: "Forbidden: No permission to view collection", } response.RespondWithError(c, response.StatusForbidden, errorResp) return } - wheres = append(wheres, model.QueryWhere{Column: "collection_id", Value: collectionID}) - + filter.CollectionID = collectionID } else { // Filter by authorization constraints - wheres = append(wheres, AuthFilter(c, authInfo)...) + filter.AuthFilters = AuthFilter(c, authInfo) } // Filter by status (support multiple values separated by comma) if statusParam := strings.TrimSpace(c.Query("status")); statusParam != "" { statusList := strings.Split(statusParam, ",") - var statusValues []interface{} + var statusValues []string 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 = statusValues } // Filter by status_not (exclude specific statuses) if statusNotParam := strings.TrimSpace(c.Query("status_not")); statusNotParam != "" { statusNotList := strings.Split(statusNotParam, ",") - var statusNotValues []interface{} + var statusNotValues []string for _, status := range statusNotList { status = strings.TrimSpace(status) if status != "" { statusNotValues = append(statusNotValues, status) } } - - if len(statusNotValues) > 0 { - if len(statusNotValues) == 1 { - // Single status exclusion - wheres = append(wheres, model.QueryWhere{ - Column: "status", - Value: statusNotValues[0], - OP: "!=", - }) - } else { - // Multiple status exclusion - use NOT IN clause - // Since gou/model doesn't support "notin" OP directly, - // we need to use a different approach or multiple != conditions - for _, status := range statusNotValues { - wheres = append(wheres, model.QueryWhere{ - Column: "status", - Value: status, - OP: "!=", - }) - } - } - } + filter.StatusNot = statusNotValues } - param.Wheres = wheres - // Add ordering sortParam := strings.TrimSpace(c.Query("sort")) if sortParam == "" { @@ -263,7 +152,7 @@ func ListDocuments(c *gin.Context) { } // Validate sort field - if !validSortFields[sortField] { + if !kbapi.ValidDocumentSortFields[sortField] { continue // Skip invalid fields } @@ -284,11 +173,10 @@ func ListDocuments(c *gin.Context) { {Column: "created_at", Option: "desc"}, } } + filter.Sort = orders - param.Orders = orders - - // Query documents using KB config - result, err := config.SearchDocuments(param, page, pagesize) + // Query documents using kb.API + result, err := kb.API.ListDocuments(c.Request.Context(), filter) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, @@ -303,8 +191,8 @@ func ListDocuments(c *gin.Context) { // GetDocument gets document details by document ID func GetDocument(c *gin.Context) { - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -318,35 +206,31 @@ func GetDocument(c *gin.Context) { return } - // Get KB instance and config - kbInstance := kb.Instance.(*kb.KnowledgeBase) - config := kbInstance.Config - // Parse select parameter - same logic as ListDocuments 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 != "" && availableDocumentFields[field] { + if field != "" && kbapi.AvailableDocumentFields[field] { selectFields = append(selectFields, field) } } // If no valid fields found, use default if len(selectFields) == 0 { - selectFields = defaultDocumentFields + selectFields = kbapi.DefaultDocumentFields } } else { - selectFields = defaultDocumentFields + selectFields = kbapi.DefaultDocumentFields } - // Build query parameters - param := model.QueryParam{ + // Build params for kb.API + params := &kbapi.GetDocumentParams{ Select: selectFields, } - // Query single document using KB config - result, err := config.FindDocument(docID, param) + // Query single document using kb.API + result, err := kb.API.GetDocument(c.Request.Context(), docID, params) if err != nil { if strings.Contains(err.Error(), "document not found") { errorResp := &response.ErrorResponse{ @@ -370,8 +254,8 @@ func GetDocument(c *gin.Context) { // RemoveDocs removes documents by IDs func RemoveDocs(c *gin.Context) { - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Check if kb.API is available + if !checkKBAPI(c) { return } @@ -405,29 +289,20 @@ func RemoveDocs(c *gin.Context) { return } - // Get KB config for database operations - 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 - } - // Validate document permissions - collectionIDs := []string{} authInfo := authorized.GetInfo(c) + checkedCollections := make(map[string]bool) for _, docID := range validDocIDs { - collectionID, _ := kbutils.ExtractCollectionIDFromDocID(docID) + collectionID := extractCollectionIDFromDocID(docID) if collectionID == "" { collectionID = "default" } - collectionIDs = append(collectionIDs, collectionID) - } - for _, collectionID := range collectionIDs { + // Skip if already checked + if checkedCollections[collectionID] { + continue + } + checkedCollections[collectionID] = true // Check update permission hasPermission, err := checkCollectionPermission(authInfo, collectionID) @@ -451,8 +326,10 @@ func RemoveDocs(c *gin.Context) { } } - // Remove documents using GraphRAG - deletedCount, err := kb.Instance.RemoveDocs(c.Request.Context(), validDocIDs) + // Remove documents using kb.API + result, err := kb.API.RemoveDocuments(c.Request.Context(), &kbapi.RemoveDocumentsParams{ + DocumentIDs: validDocIDs, + }) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, @@ -462,62 +339,16 @@ func RemoveDocs(c *gin.Context) { return } - // Also remove documents from the database and track collections to update - dbDeletedCount := 0 - collectionsToUpdate := make(map[string]bool) // Track unique collection IDs - - for _, docID := range validDocIDs { - // Get document info before deletion to track collection - if docInfo, err := config.FindDocument(docID, model.QueryParam{ - Select: []interface{}{"collection_id"}, - }); err == nil && docInfo != nil { - if collectionID, ok := docInfo["collection_id"].(string); ok && collectionID != "" { - collectionsToUpdate[collectionID] = true - } - } - - if err := config.RemoveDocument(docID); err != nil { - // Log the error but don't fail the entire operation - // since the document was already removed from GraphRAG - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to remove document from database: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return - } - dbDeletedCount++ - } - - // Update document counts for affected collections and sync to GraphRag - for collectionID := range collectionsToUpdate { - if err := UpdateDocumentCountWithSync(collectionID, config); err != nil { - // Log error but don't fail the operation - // TODO: Add proper logging - // log.Error("Failed to update document count for collection %s: %v", collectionID, err) - } - } - // Return success response with deletion count - c.JSON(http.StatusOK, gin.H{ - "message": "Documents removed successfully", - "deleted_count": deletedCount, - "requested_count": len(validDocIDs), - "db_deleted_count": dbDeletedCount, - }) + c.JSON(http.StatusOK, result) } -// Validator interface for request validation -type Validator interface { - Validate() error -} - -// checkKBInstance checks if kb.Instance is available -func checkKBInstance(c *gin.Context) bool { - if kb.Instance == nil { +// checkKBAPI checks if kb.API is available +func checkKBAPI(c *gin.Context) bool { + if kb.API == nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Knowledge base not initialized", + ErrorDescription: "Knowledge base API not initialized", } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return false @@ -525,54 +356,19 @@ func checkKBInstance(c *gin.Context) bool { return true } -// getUpsertOptions converts BaseUpsertRequest to UpsertOptions with optional file info -func getUpsertOptions(c *gin.Context, req *BaseUpsertRequest, fileInfo ...string) (*types.UpsertOptions, error) { - upsertOptions, err := req.ToUpsertOptions(fileInfo...) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: "Failed to convert request to upsert options: " + err.Error(), - } - response.RespondWithError(c, response.StatusBadRequest, errorResp) - return nil, err +// extractCollectionIDFromDocID extracts collection ID from document ID +// Document ID format: {prefix}_{collection_id}__{random_id} +func extractCollectionIDFromDocID(docID string) string { + parts := strings.Split(docID, "__") + if len(parts) < 2 { + return "" } - return upsertOptions, nil -} - -// validateFileAndGetPath validates file manager, file existence and gets local path -func validateFileAndGetPath(c *gin.Context, req *AddFileRequest) (string, string, error) { - // Get file manager - m, ok := attachment.Managers[req.Uploader] - if !ok { - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: "Invalid uploader: " + req.Uploader + " not found", - } - response.RespondWithError(c, response.StatusNotFound, errorResp) - return "", "", response.ErrInvalidRequest - } - - // Check if the file exists - exists := m.Exists(c.Request.Context(), req.FileID) - if !exists { - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: "File not found: " + req.FileID, - } - response.RespondWithError(c, response.StatusNotFound, errorResp) - return "", "", response.ErrInvalidRequest - } - - // Get the options of the manager - path, contentType, err := m.LocalPath(c.Request.Context(), req.FileID) - if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to get local path: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return "", "", err - } - - return path, contentType, nil + // First part contains prefix_collection_id + prefix := parts[0] + // Find the first underscore to skip the prefix + idx := strings.Index(prefix, "_") + if idx == -1 { + return prefix + } + return prefix[idx+1:] } diff --git a/openapi/kb/utils.go b/openapi/kb/utils.go index dc51935c..f3b182e9 100644 --- a/openapi/kb/utils.go +++ b/openapi/kb/utils.go @@ -5,9 +5,7 @@ import ( "fmt" "github.com/gin-gonic/gin" - "github.com/yaoapp/gou/graphrag/utils" "github.com/yaoapp/kun/maps" - "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" kbtypes "github.com/yaoapp/yao/kb/types" apiutils "github.com/yaoapp/yao/openapi/utils" @@ -96,134 +94,9 @@ func PrepareCreateCollection(c *gin.Context) (*CreateCollectionRequest, map[stri data["num_probes"] = req.Config.NumProbes } - // Add context fields (permissions, user info, etc.) - addContextFields(c, data) - return &req, data, nil } -// PrepareAddFile prepares AddFile request and database data -func PrepareAddFile(c *gin.Context, req *AddFileRequest) (*AddFileRequest, map[string]interface{}, error) { - // Validate request - if err := req.Validate(); err != nil { - return nil, nil, err - } - - // Validate file and get path - path, contentType, err := validateFileAndGetPath(c, req) - if err != nil { - return nil, nil, err - } - - // Get file info - m, _ := attachment.Managers[req.Uploader] - fileInfo, _ := m.Info(c.Request.Context(), req.FileID) - - // Generate document ID if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) - } - - // Prepare document data for database - data := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": fileInfo.Filename, - "type": "file", - "status": "pending", - "uploader_id": req.Uploader, - "file_name": fileInfo.Filename, - "file_path": path, - "file_mime_type": contentType, - "size": int64(fileInfo.Bytes), - } - - req.BaseUpsertRequest.AddBaseFields(data) - addContextFields(c, data) - - return req, data, nil -} - -// PrepareAddText prepares AddText request and database data -func PrepareAddText(c *gin.Context, req *AddTextRequest) (*AddTextRequest, map[string]interface{}, error) { - // Validate request - if err := req.Validate(); err != nil { - return nil, nil, err - } - - // Generate document ID if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) - } - - // Prepare document data for database - data := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": "Text Document", - "type": "text", - "status": "pending", - "text_content": req.Text, - "size": int64(len(req.Text)), - } - - // Use title from metadata if available - if req.Metadata != nil { - if title, ok := req.Metadata["title"].(string); ok && title != "" { - data["name"] = title - } - } - - req.BaseUpsertRequest.AddBaseFields(data) - addContextFields(c, data) - - return req, data, nil -} - -// PrepareAddURL prepares AddURL request and database data -func PrepareAddURL(c *gin.Context, req *AddURLRequest) (*AddURLRequest, map[string]interface{}, error) { - // Validate request - if err := req.Validate(); err != nil { - return nil, nil, err - } - - // Generate document ID if not provided - if req.DocID == "" { - req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) - } - - // Prepare document data for database - data := map[string]interface{}{ - "document_id": req.DocID, - "collection_id": req.CollectionID, - "name": req.URL, - "type": "url", - "status": "pending", - "url": req.URL, - } - - // Use title from metadata if available - if req.Metadata != nil { - if title, ok := req.Metadata["title"].(string); ok && title != "" { - data["name"] = title - data["url_title"] = title - } - } - - req.BaseUpsertRequest.AddBaseFields(data) - addContextFields(c, data) - - return req, data, nil -} - -// addContextFields adds context-specific fields like permissions, user info -func addContextFields(c *gin.Context, data map[string]interface{}) { - // TODO: Add permission-related fields from Guard - // Example: data["user_id"] = c.GetString("user_id") - // Example: data["permissions"] = c.Get("permissions") - // Example: data["tenant_id"] = c.GetString("tenant_id") -} - // UpdateCollectionWithSync updates collection metadata in database and syncs to GraphRag func UpdateCollectionWithSync(collectionID string, data maps.MapStrAny, config *kbtypes.Config) error { // Create a copy of data for GraphRag to avoid contamination from database operations From 324ee7dc92f9e6f02732023c64ea5a73949d7d36 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 15:08:16 +0800 Subject: [PATCH 03/10] Refactor Job System for Enhanced Function Execution and Error Handling - Improved the `Job` struct to support dynamic function execution with the `AddFunc` method, allowing for flexible job management. - Enhanced the `Goroutine` struct with the `ExecuteFunc` method to manage function execution, including robust error handling and context management. - Updated unit tests for `AddFunc` to ensure proper function registration and execution, including memory cleanup verification. - Revised documentation to reflect the new function execution capabilities within the job system. --- openapi/tests/kb/addfile_test.go | 368 +++++++++++++++++++++++++++ openapi/tests/kb/addtext_test.go | 401 ++++++++++++++++++++++++++++++ openapi/tests/kb/addurl_test.go | 401 ++++++++++++++++++++++++++++++ openapi/tests/kb/document_test.go | 295 ++++++++++++++++++++++ 4 files changed, 1465 insertions(+) create mode 100644 openapi/tests/kb/addfile_test.go create mode 100644 openapi/tests/kb/addtext_test.go create mode 100644 openapi/tests/kb/addurl_test.go create mode 100644 openapi/tests/kb/document_test.go diff --git a/openapi/tests/kb/addfile_test.go b/openapi/tests/kb/addfile_test.go new file mode 100644 index 00000000..1273065d --- /dev/null +++ b/openapi/tests/kb/addfile_test.go @@ -0,0 +1,368 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// TestAddFile tests the add file endpoint (sync) +func TestAddFile(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddFile Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addfile_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddFile", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddFileInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing file_id, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddFileMissingFileID", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing file_id + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains FileID (case insensitive check) + assert.Contains(t, response["error_description"], "FileID") + }) + + t.Run("AddFileNonExistentCollection", func(t *testing.T) { + // Test with a non-existent collection + addData := map[string]interface{}{ + "collection_id": "non_existent_collection_12345", + "file_id": "test_file_123", + "uploader": "local", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/non_existent_collection_12345/documents/file", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 403 Forbidden, 404 Not Found, or 500 Internal Server Error for non-existent collection + // (depends on whether permission check or collection lookup happens first) + assert.True(t, resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusInternalServerError, + "Expected 403, 404, or 500, got %d", resp.StatusCode) + }) + + t.Run("AddFileUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "file_id": "test_file_123", + "uploader": "local", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} + +// TestAddFileAsync tests the add file async endpoint +func TestAddFileAsync(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddFileAsync Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addfile_async_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddFileAsync", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddFileAsyncInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing file_id, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddFileAsyncMissingFileID", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing file_id + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains FileID (case insensitive check) + assert.Contains(t, response["error_description"], "FileID") + }) + + t.Run("AddFileAsyncUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "file_id": "test_file_123", + "uploader": "local", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) + + t.Run("AddFileAsyncFileNotFound", func(t *testing.T) { + // Test with a file_id that doesn't exist + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "file_id": "non_existent_file_12345", + "uploader": "local", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/file/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 404 Not Found for non-existent file + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + }) +} diff --git a/openapi/tests/kb/addtext_test.go b/openapi/tests/kb/addtext_test.go new file mode 100644 index 00000000..f93e250a --- /dev/null +++ b/openapi/tests/kb/addtext_test.go @@ -0,0 +1,401 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// TestAddText tests the add text endpoint (sync) +func TestAddText(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddText Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addtext_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddText", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddTextInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing text, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddTextMissingText", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing text + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Text (case insensitive check) + assert.Contains(t, response["error_description"], "Text") + }) + + t.Run("AddTextNonExistentCollection", func(t *testing.T) { + // Test with a non-existent collection + addData := map[string]interface{}{ + "collection_id": "non_existent_collection_12345", + "text": "This is a test text content for the knowledge base.", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/non_existent_collection_12345/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 403 Forbidden, 404 Not Found, or 500 Internal Server Error for non-existent collection + // (depends on whether permission check or collection lookup happens first) + assert.True(t, resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusInternalServerError, + "Expected 403, 404, or 500, got %d", resp.StatusCode) + }) + + t.Run("AddTextMissingChunking", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "text": "This is a test text content.", + // Missing chunking + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Chunking (case insensitive check) + assert.Contains(t, response["error_description"], "Chunking") + }) + + t.Run("AddTextMissingEmbedding", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "text": "This is a test text content.", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + // Missing embedding + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Embedding (case insensitive check) + assert.Contains(t, response["error_description"], "Embedding") + }) + + t.Run("AddTextUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "text": "This is a test text content.", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} + +// TestAddTextAsync tests the add text async endpoint +func TestAddTextAsync(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddTextAsync Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addtext_async_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddTextAsync", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddTextAsyncInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing text, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddTextAsyncMissingText", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing text + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Text (case insensitive check) + assert.Contains(t, response["error_description"], "Text") + }) + + t.Run("AddTextAsyncUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "text": "This is a test text content for async processing.", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/text/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} diff --git a/openapi/tests/kb/addurl_test.go b/openapi/tests/kb/addurl_test.go new file mode 100644 index 00000000..2c6ca6ab --- /dev/null +++ b/openapi/tests/kb/addurl_test.go @@ -0,0 +1,401 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// TestAddURL tests the add URL endpoint (sync) +func TestAddURL(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddURL Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addurl_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddURL", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddURLInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing url, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddURLMissingURL", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing url + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains URL (case insensitive check) + assert.Contains(t, response["error_description"], "URL") + }) + + t.Run("AddURLNonExistentCollection", func(t *testing.T) { + // Test with a non-existent collection + addData := map[string]interface{}{ + "collection_id": "non_existent_collection_12345", + "url": "https://example.com/test-page", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/non_existent_collection_12345/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 403 Forbidden, 404 Not Found, or 500 Internal Server Error for non-existent collection + // (depends on whether permission check or collection lookup happens first) + assert.True(t, resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusInternalServerError, + "Expected 403, 404, or 500, got %d", resp.StatusCode) + }) + + t.Run("AddURLMissingChunking", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "url": "https://example.com/test-page", + // Missing chunking + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Chunking (case insensitive check) + assert.Contains(t, response["error_description"], "Chunking") + }) + + t.Run("AddURLMissingEmbedding", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "url": "https://example.com/test-page", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + // Missing embedding + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains Embedding (case insensitive check) + assert.Contains(t, response["error_description"], "Embedding") + }) + + t.Run("AddURLUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "url": "https://example.com/test-page", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} + +// TestAddURLAsync tests the add URL async endpoint +func TestAddURLAsync(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB AddURLAsync Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_addurl_async_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for AddURLAsync", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("AddURLAsyncInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing url, chunking, embedding + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("AddURLAsyncMissingURL", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + // Missing url + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + // Error message contains URL (case insensitive check) + assert.Contains(t, response["error_description"], "URL") + }) + + t.Run("AddURLAsyncUnauthorized", func(t *testing.T) { + addData := map[string]interface{}{ + "collection_id": testCollectionID, + "url": "https://example.com/async-test-page", + "chunking": map[string]interface{}{ + "provider_id": "__yao.structured", + "option_id": "standard", + }, + "embedding": map[string]interface{}{ + "provider_id": "__yao.openai", + "option_id": "text-embedding-3-small", + }, + } + + body, err := json.Marshal(addData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/documents/url/async", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} diff --git a/openapi/tests/kb/document_test.go b/openapi/tests/kb/document_test.go new file mode 100644 index 00000000..01fe937a --- /dev/null +++ b/openapi/tests/kb/document_test.go @@ -0,0 +1,295 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// TestListDocuments tests the document listing endpoint +func TestListDocuments(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB Document List Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Create a test collection first + testCollectionID := fmt.Sprintf("test_doc_list_collection_%d", time.Now().UnixNano()) + testutils.RegisterTestCollection(testCollectionID) + + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "name": "Test Collection for Document List", + "category": "test", + }, + "config": map[string]interface{}{ + "embedding_provider_id": "__yao.openai", + "embedding_option_id": "text-embedding-3-small", + "locale": "en", + "index_type": "hnsw", + "distance": "cosine", + }, + } + + body, _ := json.Marshal(createData) + req, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("Failed to create test collection: %v", err) + } + resp.Body.Close() + + t.Run("ListDocumentsSuccess", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents?collection_id="+testCollectionID, nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + + // Verify pagination fields exist + assert.Contains(t, response, "data") + assert.Contains(t, response, "page") + assert.Contains(t, response, "pagesize") + assert.Contains(t, response, "total") + }) + + t.Run("ListDocumentsWithPagination", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents?collection_id="+testCollectionID+"&page=1&pagesize=10", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + + // Verify pagination values + assert.Equal(t, float64(1), response["page"]) + assert.Equal(t, float64(10), response["pagesize"]) + }) + + t.Run("ListDocumentsWithStatusFilter", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents?collection_id="+testCollectionID+"&status=completed", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + }) + + t.Run("ListDocumentsWithSort", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents?collection_id="+testCollectionID+"&sort=created_at+desc", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + }) + + t.Run("ListDocumentsUnauthorized", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents?collection_id="+testCollectionID, nil) + assert.NoError(t, err) + + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} + +// TestGetDocument tests the get document endpoint +func TestGetDocument(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB Document Get Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + t.Run("GetDocumentNotFound", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents/non_existent_doc_id", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 404 Not Found + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + }) + + t.Run("GetDocumentWithSelectFields", func(t *testing.T) { + // This test verifies the select parameter works + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents/test_doc_id?select=id,name,status", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 404 since doc doesn't exist, but the request format is valid + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + }) + + t.Run("GetDocumentUnauthorized", func(t *testing.T) { + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/documents/test_doc_id", nil) + assert.NoError(t, err) + + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} + +// TestRemoveDocuments tests the remove documents endpoint +func TestRemoveDocuments(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB Document Remove Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + t.Run("RemoveDocumentsMissingIDs", func(t *testing.T) { + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/documents", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request for missing document_ids + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + assert.Contains(t, response["error_description"], "document_ids") + }) + + t.Run("RemoveDocumentsEmptyIDs", func(t *testing.T) { + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/documents?document_ids=", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 400 Bad Request for empty document_ids + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + }) + + t.Run("RemoveDocumentsNonExistent", func(t *testing.T) { + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/documents?document_ids=non_existent_doc_1,non_existent_doc_2", nil) + assert.NoError(t, err) + + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // The behavior depends on implementation - could be 200 with 0 removed or 404 + // Accept either as valid + assert.True(t, resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusForbidden) + }) + + t.Run("RemoveDocumentsUnauthorized", func(t *testing.T) { + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/documents?document_ids=doc1,doc2", nil) + assert.NoError(t, err) + + // No Authorization header + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Should return 401 Unauthorized + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + }) +} From 8dcb78fc3793f48ddf70b4a61d9d892cb712bc91 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 15:26:18 +0800 Subject: [PATCH 04/10] Update serpAPIKnowledge struct to allow flexible type for 'Type' field - Modified the 'Type' field in the serpAPIKnowledge struct to use an interface{}, enabling it to accept either a string or an object based on the query context. - This change enhances the flexibility of the knowledge graph data representation in the search handler. --- agent/search/handlers/web/serpapi.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/agent/search/handlers/web/serpapi.go b/agent/search/handlers/web/serpapi.go index 623860a8..3e9cf5af 100644 --- a/agent/search/handlers/web/serpapi.go +++ b/agent/search/handlers/web/serpapi.go @@ -115,9 +115,9 @@ type serpAPIAnswerBox struct { // serpAPIKnowledge represents knowledge graph data type serpAPIKnowledge struct { - Title string `json:"title,omitempty"` - Type string `json:"type,omitempty"` - Description string `json:"description,omitempty"` + Title string `json:"title,omitempty"` + Type interface{} `json:"type,omitempty"` // Can be string or object depending on query + Description string `json:"description,omitempty"` } // serpAPIRelated represents related searches From 493f9dbea2373c2603a611ad4045144fe47decd6 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 19:48:00 +0800 Subject: [PATCH 05/10] Update test case for SerpAPIProvider to reflect new query and expected result - Changed the query in the `TestSerpAPIProviderWithAssistantConfig` from "Yao App Engine" to "golang programming language" to align with updated test expectations. - Updated the expected result assertion to match the new query, ensuring the test accurately verifies the functionality of the search handler. --- agent/search/handlers/web/serpapi_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/agent/search/handlers/web/serpapi_test.go b/agent/search/handlers/web/serpapi_test.go index a4637e5d..0a6ae6b1 100644 --- a/agent/search/handlers/web/serpapi_test.go +++ b/agent/search/handlers/web/serpapi_test.go @@ -39,7 +39,7 @@ func TestSerpAPIProviderWithAssistantConfig(t *testing.T) { // Execute search req := &types.Request{ - Query: "Yao App Engine", + Query: "golang programming language", Type: types.SearchTypeWeb, Source: types.SourceAuto, Limit: 5, @@ -51,7 +51,7 @@ func TestSerpAPIProviderWithAssistantConfig(t *testing.T) { // Verify result structure assert.Equal(t, types.SearchTypeWeb, result.Type) - assert.Equal(t, "Yao App Engine", result.Query) + assert.Equal(t, "golang programming language", result.Query) assert.Equal(t, types.SourceAuto, result.Source) // API key must be valid - search should succeed From fa471bc4c415afe8c3b2f0283b5b572e7a39db06 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 19 Dec 2025 20:10:16 +0800 Subject: [PATCH 06/10] Update tests to use context and correct URL references - Modified the `createTestContext` function to use `context.Background()` for improved context management in tests. - Updated the URL in the `TestAddURL` function to point to the correct Yao Agent Caller resource, ensuring accurate test assertions for added URLs. --- agent/search/handlers/web/agent_test.go | 3 ++- kb/api/addurl_test.go | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/agent/search/handlers/web/agent_test.go b/agent/search/handlers/web/agent_test.go index 8438c06f..d6fafbdf 100644 --- a/agent/search/handlers/web/agent_test.go +++ b/agent/search/handlers/web/agent_test.go @@ -1,6 +1,7 @@ package web_test import ( + "context" "testing" "github.com/stretchr/testify/assert" @@ -278,7 +279,7 @@ func createTestContext(t *testing.T) *agentContext.Context { UserID: "test-user", TenantID: "test-tenant", } - ctx := agentContext.New(nil, authorized, "test-chat-id") + ctx := agentContext.New(context.Background(), authorized, "test-chat-id") ctx.AssistantID = "tests.web-agent-caller" return ctx } diff --git a/kb/api/addurl_test.go b/kb/api/addurl_test.go index fc48827f..41d5db15 100644 --- a/kb/api/addurl_test.go +++ b/kb/api/addurl_test.go @@ -67,10 +67,10 @@ func TestAddURL(t *testing.T) { t.Run("AddURLSuccess", func(t *testing.T) { params := &api.AddURLParams{ CollectionID: collectionID, - URL: "https://example.com", + URL: "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go", Locale: "en", Metadata: map[string]interface{}{ - "title": "Example Website", + "title": "Yao Agent Caller", "description": "A test URL document", }, Chunking: &api.ProviderConfigParams{ @@ -96,7 +96,7 @@ func TestAddURL(t *testing.T) { if result != nil { assert.Equal(t, collectionID, result.CollectionID) assert.NotEmpty(t, result.DocID) - assert.Equal(t, "https://example.com", result.URL) + assert.Equal(t, "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go", result.URL) assert.Contains(t, result.Message, "successfully") t.Logf("Added URL document: %s", result.DocID) } @@ -107,7 +107,7 @@ func TestAddURL(t *testing.T) { assert.NoError(t, err) assert.NotNil(t, doc) assert.Equal(t, "url", doc["type"]) - assert.Equal(t, "https://example.com", doc["url"]) + assert.Equal(t, "https://raw.githubusercontent.com/trheyi/yao/refs/heads/main/agent/caller/caller.go", doc["url"]) t.Logf("✅ URL Document verified: type=%v, url=%v, status=%v", doc["type"], doc["url"], doc["status"]) } }) From 53db8be522c1aff402355dabd18e63232ed93dd4 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 20 Dec 2025 11:25:02 +0800 Subject: [PATCH 07/10] Enhance Collection Retrieval and Existence Check Logic - Updated the `GetCollection` method to first read from the database for existence and permissions, improving data integrity. - Merged metadata from GraphRag into the result, ensuring backward compatibility and enhanced data representation. - Refactored the `CollectionExists` method to check both the database and GraphRag for consistency, logging any mismatches for debugging purposes. - Introduced new types and structures for search operations, including `SearchMode`, `Query`, and `SearchResult`, to support advanced search functionalities. --- kb/api/collection.go | 72 +++-- kb/api/interfaces.go | 4 +- kb/api/search.go | 551 ++++++++++++++++++++++++++++++++++++ kb/api/search_setup_test.go | 400 ++++++++++++++++++++++++++ kb/api/search_test.go | 446 +++++++++++++++++++++++++++++ kb/api/types.go | 77 +++++ 6 files changed, 1521 insertions(+), 29 deletions(-) create mode 100644 kb/api/search.go create mode 100644 kb/api/search_setup_test.go create mode 100644 kb/api/search_test.go diff --git a/kb/api/collection.go b/kb/api/collection.go index 051ea758..916ce800 100644 --- a/kb/api/collection.go +++ b/kb/api/collection.go @@ -274,61 +274,79 @@ func (instance *KBInstance) RemoveCollection(ctx context.Context, collectionID s } // GetCollection retrieves a collection by ID +// Reads from database first, then merges with GraphRag metadata 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) + // Read from database (source of truth for existence and permissions) + dbRecord, err := instance.Config.FindCollection(collectionID, model.QueryParam{}) 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) + return nil, fmt.Errorf("collection not found") } - // 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) + // Convert database record to result map (flatten to top level) 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 + for k, v := range dbRecord { + result[k] = v } - if collection.Config != nil { - result["config"] = collection.Config + // Set standard ID fields + result["id"] = collectionID + result["collection_id"] = collectionID + + // Read from GraphRag and merge (for config and metadata object) + graphRagCollection, err := instance.GraphRag.GetCollection(ctx, collectionID) + if err == nil && graphRagCollection != nil { + // Set GraphRag config (vector store configuration) + if graphRagCollection.Config != nil { + result["config"] = graphRagCollection.Config + } + + // Set GraphRag metadata as nested object (for backward compatibility) + // This allows access via collection["metadata"]["field"] + if graphRagCollection.Metadata != nil { + result["metadata"] = graphRagCollection.Metadata + + // Also flatten GraphRag metadata fields to top level + // Only add fields that don't exist in database record + for k, v := range graphRagCollection.Metadata { + if _, exists := result[k]; !exists { + result[k] = v + } + } + } } return result, nil } // CollectionExists checks if a collection exists by ID +// Checks both database and GraphRag for consistency 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) + // Check database (source of truth for existence) + _, dbErr := instance.Config.FindCollection(collectionID, model.QueryParam{}) + dbExists := dbErr == nil + + // Check GraphRag for consistency + graphRagExists, _ := instance.GraphRag.CollectionExists(ctx, collectionID) + + // Collection exists if it exists in database + // Log warning if there's inconsistency (for debugging) + if dbExists != graphRagExists { + log.Warn("Collection %s existence mismatch: database=%v, graphrag=%v", collectionID, dbExists, graphRagExists) } return &CollectionExistsResult{ CollectionID: collectionID, - Exists: exists, + Exists: dbExists, }, nil } diff --git a/kb/api/interfaces.go b/kb/api/interfaces.go index 5ac527b7..3cab0013 100644 --- a/kb/api/interfaces.go +++ b/kb/api/interfaces.go @@ -32,8 +32,8 @@ type API interface { AddTextAsync(ctx context.Context, params *AddTextParams) (*AddDocumentAsyncResult, error) AddURLAsync(ctx context.Context, params *AddURLParams) (*AddDocumentAsyncResult, error) - // Segment operations (future) - // ... + // Search operations + Search(ctx context.Context, queries []Query) (*SearchResult, error) } // KBInstance holds the KB instance dependencies required by the API diff --git a/kb/api/search.go b/kb/api/search.go new file mode 100644 index 00000000..cc1bfe4d --- /dev/null +++ b/kb/api/search.go @@ -0,0 +1,551 @@ +package api + +import ( + "context" + "fmt" + "sort" + "sync" + + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/kb/providers/factory" +) + +// Default search parameters +const ( + DefaultSearchK = 10 + DefaultMaxDepth = 2 + DefaultMinScore = 0.0 + MaxSearchK = 100 + DefaultSearchPageSize = 20 +) + +// Search performs batch search operations on the knowledge base +// Queries can span multiple collections; implementation groups by CollectionID +// Mode, providers (embedding/extraction/reranker) are read from each collection's config +// All results are merged and deduplicated +func (kb *KBInstance) Search(ctx context.Context, queries []Query) (*SearchResult, error) { + if len(queries) == 0 { + return &SearchResult{ + Segments: []graphragtypes.Segment{}, + Total: 0, + }, nil + } + + // 1. Validate queries + if err := kb.validateQueries(queries); err != nil { + return nil, err + } + + // 2. Group queries by CollectionID + groupedQueries := kb.groupQueriesByCollection(queries) + + // 3. Process each collection group in parallel + var ( + allSegments []graphragtypes.Segment + allGraph *GraphData + mu sync.Mutex + wg sync.WaitGroup + errChan = make(chan error, len(groupedQueries)) + ) + + for collectionID, collQueries := range groupedQueries { + wg.Add(1) + go func(collID string, qs []Query) { + defer wg.Done() + + segments, graph, err := kb.searchCollection(ctx, collID, qs) + if err != nil { + errChan <- fmt.Errorf("search in collection %s failed: %w", collID, err) + return + } + + mu.Lock() + allSegments = append(allSegments, segments...) + if graph != nil { + allGraph = mergeGraphData(allGraph, graph) + } + mu.Unlock() + }(collectionID, collQueries) + } + + wg.Wait() + close(errChan) + + // Collect errors + var errors []error + for err := range errChan { + errors = append(errors, err) + } + if len(errors) > 0 { + log.Warn("Search completed with errors: %v", errors) + } + + // 4. Merge and deduplicate results + mergedSegments := kb.deduplicateSegments(allSegments) + + // 5. Sort by score (descending) + sort.Slice(mergedSegments, func(i, j int) bool { + return mergedSegments[i].Score > mergedSegments[j].Score + }) + + // 6. Apply pagination from first query (if specified) + result := kb.applyPagination(mergedSegments, queries[0]) + result.Graph = allGraph + + return result, nil +} + +// ========== Validation ========== + +// validateQueries validates all queries +func (kb *KBInstance) validateQueries(queries []Query) error { + for i, q := range queries { + if q.CollectionID == "" { + return fmt.Errorf("query %d: collection_id is required", i) + } + if q.Input == "" && len(q.Messages) == 0 { + return fmt.Errorf("query %d: either input or messages is required", i) + } + } + return nil +} + +// ========== Query Grouping ========== + +// groupQueriesByCollection groups queries by their CollectionID +func (kb *KBInstance) groupQueriesByCollection(queries []Query) map[string][]Query { + grouped := make(map[string][]Query) + for _, q := range queries { + grouped[q.CollectionID] = append(grouped[q.CollectionID], q) + } + return grouped +} + +// ========== Collection Search ========== + +// searchCollection processes all queries for a single collection +func (kb *KBInstance) searchCollection(ctx context.Context, collectionID string, queries []Query) ([]graphragtypes.Segment, *GraphData, error) { + // Get collection config + collection, err := kb.GetCollection(ctx, collectionID) + if err != nil { + return nil, nil, fmt.Errorf("failed to get collection: %w", err) + } + + // Get embedding provider from collection config + embeddingProviderID, _ := collection["embedding_provider_id"].(string) + embeddingOptionID, _ := collection["embedding_option_id"].(string) + if embeddingProviderID == "" || embeddingOptionID == "" { + return nil, nil, fmt.Errorf("collection %s missing embedding configuration", collectionID) + } + + // Create embedding function + embedding, err := kb.createEmbedding(embeddingProviderID, embeddingOptionID, "en") + if err != nil { + return nil, nil, fmt.Errorf("failed to create embedding: %w", err) + } + + var ( + allSegments []graphragtypes.Segment + allGraph *GraphData + mu sync.Mutex + wg sync.WaitGroup + errChan = make(chan error, len(queries)) + ) + + // Process queries in parallel + for _, query := range queries { + wg.Add(1) + go func(q Query) { + defer wg.Done() + + segments, graph, err := kb.executeQuery(ctx, collectionID, q, embedding, collection) + if err != nil { + errChan <- err + return + } + + mu.Lock() + allSegments = append(allSegments, segments...) + if graph != nil { + allGraph = mergeGraphData(allGraph, graph) + } + mu.Unlock() + }(query) + } + + wg.Wait() + close(errChan) + + // Collect errors + var errors []error + for err := range errChan { + errors = append(errors, err) + } + if len(errors) > 0 { + return allSegments, allGraph, errors[0] + } + + return allSegments, allGraph, nil +} + +// executeQuery executes a single query based on its mode +func (kb *KBInstance) executeQuery(ctx context.Context, collectionID string, query Query, embedding graphragtypes.Embedding, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) { + // Determine search mode + mode := query.Mode + if mode == "" { + // Default to expand mode + mode = SearchModeExpand + } + + // Get query text + queryText := kb.getQueryText(query) + if queryText == "" { + return nil, nil, fmt.Errorf("no query text found") + } + + // Execute based on mode + switch mode { + case SearchModeVector: + return kb.searchVector(ctx, collectionID, queryText, query, embedding) + case SearchModeGraph: + return kb.searchGraph(ctx, collectionID, queryText, query, collection) + case SearchModeExpand: + return kb.searchExpand(ctx, collectionID, queryText, query, embedding, collection) + default: + return nil, nil, fmt.Errorf("unknown search mode: %s", mode) + } +} + +// getQueryText extracts query text from Input or Messages +func (kb *KBInstance) getQueryText(query Query) string { + // Input takes precedence + if query.Input != "" { + return query.Input + } + + // Extract from last user message + for i := len(query.Messages) - 1; i >= 0; i-- { + if query.Messages[i].Role == "user" { + return query.Messages[i].Content + } + } + + return "" +} + +// ========== Vector Search ========== + +// searchVector performs pure vector similarity search +func (kb *KBInstance) searchVector(ctx context.Context, collectionID string, queryText string, query Query, embedding graphragtypes.Embedding) ([]graphragtypes.Segment, *GraphData, error) { + // Build search options + k := query.PageSize + if k <= 0 { + k = DefaultSearchK + } + if k > MaxSearchK { + k = MaxSearchK + } + + options := &graphragtypes.VectorSearchOptions{ + CollectionID: collectionID, + DocumentID: query.DocumentID, + Query: queryText, + K: k, + MinScore: query.MinScore, + Embedding: embedding, + } + + // Add metadata filter + if len(query.Metadata) > 0 { + options.Filter = query.Metadata + } + + // Execute search + result, err := kb.GraphRag.SearchVector(ctx, options) + if err != nil { + return nil, nil, fmt.Errorf("vector search failed: %w", err) + } + return result.Segments, nil, nil +} + +// ========== Graph Search ========== + +// searchGraph performs pure graph traversal search +func (kb *KBInstance) searchGraph(ctx context.Context, collectionID string, queryText string, query Query, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) { + // Get extraction provider for entity extraction + extraction, err := kb.createExtraction(collection) + if err != nil { + return nil, nil, fmt.Errorf("failed to create extraction: %w", err) + } + + // Build graph search options + maxDepth := query.MaxDepth + if maxDepth <= 0 { + maxDepth = DefaultMaxDepth + } + + options := &graphragtypes.GraphSearchOptions{ + CollectionID: collectionID, + DocumentID: query.DocumentID, + Query: queryText, + MaxDepth: maxDepth, + Extraction: extraction, + } + + // Execute search + result, err := kb.GraphRag.SearchGraph(ctx, options) + if err != nil { + return nil, nil, fmt.Errorf("graph search failed: %w", err) + } + + // Convert to GraphData + graph := &GraphData{ + Nodes: result.Nodes, + Relationships: result.Relationships, + } + + return result.Segments, graph, nil +} + +// ========== Expand Search (Graph + Vector) ========== + +// searchExpand performs graph-based entity expansion + vector search +// This mode uses graph to find related entities, then enhances vector search +func (kb *KBInstance) searchExpand(ctx context.Context, collectionID string, queryText string, query Query, embedding graphragtypes.Embedding, collection map[string]interface{}) ([]graphragtypes.Segment, *GraphData, error) { + // Step 1: Extract entities from query using graph search + extraction, err := kb.createExtraction(collection) + if err != nil { + // Fall back to pure vector search if extraction is not available + log.Warn("Extraction not available, falling back to vector search: %v", err) + return kb.searchVector(ctx, collectionID, queryText, query, embedding) + } + + maxDepth := query.MaxDepth + if maxDepth <= 0 { + maxDepth = DefaultMaxDepth + } + + graphOptions := &graphragtypes.GraphSearchOptions{ + CollectionID: collectionID, + DocumentID: query.DocumentID, + Query: queryText, + MaxDepth: maxDepth, + Extraction: extraction, + } + + // Execute graph search to find related entities + graphResult, graphErr := kb.GraphRag.SearchGraph(ctx, graphOptions) + + // Step 2: Perform vector search + k := query.PageSize + if k <= 0 { + k = DefaultSearchK + } + if k > MaxSearchK { + k = MaxSearchK + } + + vectorOptions := &graphragtypes.VectorSearchOptions{ + CollectionID: collectionID, + DocumentID: query.DocumentID, + Query: queryText, + K: k, + MinScore: query.MinScore, + Embedding: embedding, + } + + if len(query.Metadata) > 0 { + vectorOptions.Filter = query.Metadata + } + + vectorResult, err := kb.GraphRag.SearchVector(ctx, vectorOptions) + if err != nil { + return nil, nil, fmt.Errorf("vector search failed: %w", err) + } + + // Step 3: Merge results + segments := vectorResult.Segments + + var graph *GraphData + if graphErr == nil && graphResult != nil { + // Add graph segments (deduplicated later) + segments = append(segments, graphResult.Segments...) + + // Include graph data + graph = &GraphData{ + Nodes: graphResult.Nodes, + Relationships: graphResult.Relationships, + } + } + + return segments, graph, nil +} + +// ========== Helper Functions ========== + +// createEmbedding creates an embedding function from provider config +func (kb *KBInstance) createEmbedding(providerID, optionID, locale string) (graphragtypes.Embedding, error) { + if locale == "" { + locale = "en" + } + + // Get provider option + option, err := kb.getProviderOption("embedding", providerID, optionID, locale) + if err != nil { + return nil, fmt.Errorf("failed to get embedding option: %w", err) + } + + // Create embedding provider + return factory.MakeEmbedding(providerID, option) +} + +// createExtraction creates an extraction function from collection config +func (kb *KBInstance) createExtraction(collection map[string]interface{}) (graphragtypes.Extraction, error) { + // Try to get extraction provider from collection metadata + metadata, _ := collection["metadata"].(map[string]interface{}) + if metadata == nil { + metadata = collection + } + + extractionProviderID, _ := metadata["__extraction_provider"].(string) + extractionOptionID, _ := metadata["__extraction_option"].(string) + + // Fall back to default extraction provider + if extractionProviderID == "" { + extractionProviderID = "__yao.openai" + extractionOptionID = "gpt-4o-mini" + } + + // Get provider option + option, err := kb.getProviderOption("extraction", extractionProviderID, extractionOptionID, "en") + if err != nil { + return nil, fmt.Errorf("failed to get extraction option: %w", err) + } + + // Create extraction provider + return factory.MakeExtraction(extractionProviderID, option) +} + +// deduplicateSegments removes duplicate segments by ID, keeping highest score +func (kb *KBInstance) deduplicateSegments(segments []graphragtypes.Segment) []graphragtypes.Segment { + seen := make(map[string]int) // ID -> index in result + result := make([]graphragtypes.Segment, 0, len(segments)) + + for _, seg := range segments { + if idx, exists := seen[seg.ID]; exists { + // Keep the one with higher score + if seg.Score > result[idx].Score { + result[idx] = seg + } + } else { + seen[seg.ID] = len(result) + result = append(result, seg) + } + } + + return result +} + +// mergeGraphData merges two GraphData objects +func mergeGraphData(a, b *GraphData) *GraphData { + if a == nil { + return b + } + if b == nil { + return a + } + + // Merge nodes (deduplicate by ID) + nodeMap := make(map[string]graphragtypes.GraphNode) + for _, n := range a.Nodes { + nodeMap[n.ID] = n + } + for _, n := range b.Nodes { + nodeMap[n.ID] = n + } + + nodes := make([]graphragtypes.GraphNode, 0, len(nodeMap)) + for _, n := range nodeMap { + nodes = append(nodes, n) + } + + // Merge relationships (deduplicate by ID) + relMap := make(map[string]graphragtypes.GraphRelationship) + for _, r := range a.Relationships { + relMap[r.ID] = r + } + for _, r := range b.Relationships { + relMap[r.ID] = r + } + + relationships := make([]graphragtypes.GraphRelationship, 0, len(relMap)) + for _, r := range relMap { + relationships = append(relationships, r) + } + + return &GraphData{ + Nodes: nodes, + Relationships: relationships, + } +} + +// applyPagination applies pagination to segments +func (kb *KBInstance) applyPagination(segments []graphragtypes.Segment, query Query) *SearchResult { + total := len(segments) + + // If no pagination requested, return all + if query.Page <= 0 && query.PageSize <= 0 { + return &SearchResult{ + Segments: segments, + Total: total, + } + } + + page := query.Page + if page <= 0 { + page = 1 + } + + pageSize := query.PageSize + if pageSize <= 0 { + pageSize = DefaultSearchPageSize + } + + // Calculate pagination + totalPages := (total + pageSize - 1) / pageSize + start := (page - 1) * pageSize + end := start + pageSize + + if start >= total { + return &SearchResult{ + Segments: []graphragtypes.Segment{}, + Total: total, + Page: page, + PageSize: pageSize, + TotalPages: totalPages, + } + } + + if end > total { + end = total + } + + result := &SearchResult{ + Segments: segments[start:end], + Total: total, + Page: page, + PageSize: pageSize, + TotalPages: totalPages, + } + + // Set next/prev page + if page < totalPages { + result.Next = page + 1 + } + if page > 1 { + result.Prev = page - 1 + } + + return result +} diff --git a/kb/api/search_setup_test.go b/kb/api/search_setup_test.go new file mode 100644 index 00000000..000bb0cc --- /dev/null +++ b/kb/api/search_setup_test.go @@ -0,0 +1,400 @@ +package api_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go + +// ========== Fixed Test Collection IDs ========== +// Use fixed IDs so we can reuse them across test runs during development + +const ( + // SearchTestScienceCollection is the fixed ID for science test collection + SearchTestScienceCollection = "search_test_science" + // SearchTestTechCollection is the fixed ID for tech test collection + SearchTestTechCollection = "search_test_tech" +) + +// ========== Setup Test - Run Once ========== + +// TestSearchSetup creates test collections and documents for search testing. +// Run this once before running search tests: +// +// go test -v -run "TestSearchSetup" ./kb/api/... +// +// Then run search tests multiple times without waiting for data setup: +// +// go test -v -run "TestSearchQuery" ./kb/api/... +func TestSearchSetup(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + + // Check if collections already exist and are complete + // We check both GraphRag (vector store) and document count + scienceComplete := false + techComplete := false + + // Check Science collection + scienceCollection, scienceErr := kb.API.GetCollection(ctx, SearchTestScienceCollection) + if scienceErr == nil && scienceCollection != nil { + scienceDocs, _ := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestScienceCollection, + }) + if scienceDocs != nil && len(scienceDocs.Data) >= 5 { + scienceComplete = true + t.Logf("✓ Science collection exists: %s (%d docs)", SearchTestScienceCollection, len(scienceDocs.Data)) + } + } + + // Check Tech collection + techCollection, techErr := kb.API.GetCollection(ctx, SearchTestTechCollection) + if techErr == nil && techCollection != nil { + techDocs, _ := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestTechCollection, + }) + if techDocs != nil && len(techDocs.Data) >= 5 { + techComplete = true + t.Logf("✓ Tech collection exists: %s (%d docs)", SearchTestTechCollection, len(techDocs.Data)) + } + } + + // If both collections are complete, skip setup + if scienceComplete && techComplete { + t.Log("✓ All test collections already exist with sufficient documents") + t.Log(" Skipping setup. Run TestSearchCleanup first to recreate.") + return + } + + // Clean up any existing collections (handles both complete and incomplete states) + // RemoveCollection cleans both database and GraphRag (including orphaned vector collections) + t.Log("Cleaning up existing collections...") + if result, err := kb.API.RemoveCollection(ctx, SearchTestScienceCollection); err == nil && result.Removed { + t.Logf(" Removed: %s", SearchTestScienceCollection) + } + if result, err := kb.API.RemoveCollection(ctx, SearchTestTechCollection); err == nil && result.Removed { + t.Logf(" Removed: %s", SearchTestTechCollection) + } + time.Sleep(1 * time.Second) // Wait for cleanup + + // Create Science Collection + t.Log("Creating Science collection...") + scienceParams := &api.CreateCollectionParams{ + ID: SearchTestScienceCollection, + Metadata: map[string]interface{}{ + "name": "Science Knowledge Base", + "description": "Scientists and their discoveries for search testing", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + _, err := kb.API.CreateCollection(ctx, scienceParams) + if err != nil { + t.Fatalf("Failed to create science collection: %v", err) + } + t.Logf("✓ Created collection: %s", SearchTestScienceCollection) + + // Create Tech Collection + t.Log("Creating Tech collection...") + techParams := &api.CreateCollectionParams{ + ID: SearchTestTechCollection, + Metadata: map[string]interface{}{ + "name": "Tech Knowledge Base", + "description": "Technology companies and products for search testing", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &graphragtypes.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, + } + _, err = kb.API.CreateCollection(ctx, techParams) + if err != nil { + t.Fatalf("Failed to create tech collection: %v", err) + } + t.Logf("✓ Created collection: %s", SearchTestTechCollection) + + // Add Science Documents + // Entity relationships: Einstein -> Relativity -> Physics -> Nobel Prize + scienceDocs := []struct { + title string + content string + }{ + { + title: "Albert Einstein Biography", + content: `Albert Einstein was a theoretical physicist born in Germany in 1879. + He developed the theory of relativity, one of the two pillars of modern physics. + Einstein received the Nobel Prize in Physics in 1921 for his discovery of the photoelectric effect. + He later emigrated to the United States and worked at Princeton University until his death in 1955.`, + }, + { + title: "Theory of Relativity", + content: `The theory of relativity was developed by Albert Einstein in the early 20th century. + It consists of special relativity (1905) and general relativity (1915). + Special relativity introduced E=mc², showing the relationship between energy and mass. + General relativity describes gravity as the curvature of spacetime caused by mass and energy.`, + }, + { + title: "Marie Curie Biography", + content: `Marie Curie was a Polish-French physicist and chemist who conducted pioneering research on radioactivity. + She was the first woman to win a Nobel Prize and the only person to win Nobel Prizes in two different sciences (Physics and Chemistry). + Curie discovered the elements polonium and radium. She founded the Curie Institutes in Paris and Warsaw.`, + }, + { + title: "Nobel Prize in Physics", + content: `The Nobel Prize in Physics is awarded annually by the Royal Swedish Academy of Sciences. + Notable recipients include Albert Einstein (1921) for the photoelectric effect, + Marie Curie (1903) for research on radiation phenomena, + and Niels Bohr (1922) for his contributions to understanding atomic structure.`, + }, + { + title: "Quantum Mechanics Foundations", + content: `Quantum mechanics emerged in the early 20th century through the work of many physicists. + Max Planck introduced the concept of energy quanta in 1900. + Niels Bohr proposed the Bohr model of the atom. + Werner Heisenberg developed the uncertainty principle. + These discoveries built upon Einstein's work on the photoelectric effect.`, + }, + } + + t.Log("Adding Science documents...") + for _, doc := range scienceDocs { + docID := addFixedTestDocument(t, ctx, SearchTestScienceCollection, doc.title, doc.content) + if docID != "" { + t.Logf(" ✓ Added: %s", doc.title) + } + } + + // Add Tech Documents + // Entity relationships: Apple -> Steve Jobs -> iPhone -> iOS + techDocs := []struct { + title string + content string + }{ + { + title: "Apple Inc History", + content: `Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in 1976. + The company revolutionized personal computing with the Macintosh in 1984. + Under Steve Jobs' leadership, Apple introduced the iPhone in 2007, which transformed the smartphone industry. + Apple is headquartered in Cupertino, California.`, + }, + { + title: "iPhone Development", + content: `The iPhone was introduced by Steve Jobs at Macworld 2007. + It combined a mobile phone, widescreen iPod, and internet device into one product. + The iPhone runs on iOS, Apple's mobile operating system. + The App Store, launched in 2008, created a new ecosystem for mobile applications.`, + }, + { + title: "Google and AI", + content: `Google has been a pioneer in artificial intelligence and machine learning. + The company developed TensorFlow, an open-source machine learning framework. + Google's AI research includes natural language processing, computer vision, and deep learning. + Google Brain and DeepMind are the company's main AI research divisions.`, + }, + { + title: "Machine Learning Applications", + content: `Machine learning is transforming various industries through AI applications. + Google uses ML for search ranking, language translation, and image recognition. + TensorFlow enables developers to build and train neural networks. + Deep learning models can now understand natural language and generate human-like text.`, + }, + { + title: "Tech Industry Leaders", + content: `The technology industry has been shaped by visionary leaders. + Steve Jobs transformed Apple into the world's most valuable company. + Larry Page and Sergey Brin founded Google and pioneered internet search. + Elon Musk leads Tesla and SpaceX, pushing boundaries in electric vehicles and space exploration.`, + }, + } + + t.Log("Adding Tech documents...") + for _, doc := range techDocs { + docID := addFixedTestDocument(t, ctx, SearchTestTechCollection, doc.title, doc.content) + if docID != "" { + t.Logf(" ✓ Added: %s", doc.title) + } + } + + // Wait for indexing + t.Log("Waiting for indexing...") + time.Sleep(2 * time.Second) + + // Verify setup + t.Log("Verifying setup...") + scienceDocsResult, _ := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestScienceCollection, + }) + techDocsResult, _ := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestTechCollection, + }) + + t.Logf("✓ Setup complete!") + t.Logf(" Science collection: %d documents", len(scienceDocsResult.Data)) + t.Logf(" Tech collection: %d documents", len(techDocsResult.Data)) + t.Logf("") + t.Logf("Now run search tests with:") + t.Logf(" go test -v -run 'TestSearchQuery' ./kb/api/...") +} + +// ========== Cleanup Test ========== + +// TestSearchCleanup removes test collections. +// Run this to clean up test data: +// +// go test -v -run "TestSearchCleanup" ./kb/api/... +func TestSearchCleanup(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + + t.Log("Removing test collections...") + + result1, err := kb.API.RemoveCollection(ctx, SearchTestScienceCollection) + if err != nil { + t.Logf(" Science collection removal: %v", err) + } else if result1.Removed { + t.Logf("✓ Removed: %s", SearchTestScienceCollection) + } + + result2, err := kb.API.RemoveCollection(ctx, SearchTestTechCollection) + if err != nil { + t.Logf(" Tech collection removal: %v", err) + } else if result2.Removed { + t.Logf("✓ Removed: %s", SearchTestTechCollection) + } + + t.Log("✓ Cleanup complete!") +} + +// ========== Verify Test ========== + +// TestSearchVerify checks if test collections exist and have documents. +// Run this to verify test data: +// +// go test -v -run "TestSearchVerify" ./kb/api/... +func TestSearchVerify(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + + // Check Science collection + scienceExists, err := kb.API.CollectionExists(ctx, SearchTestScienceCollection) + if err != nil { + t.Fatalf("Failed to check science collection: %v", err) + } + if !scienceExists.Exists { + t.Fatalf("✗ Science collection does not exist. Run TestSearchSetup first.") + } + + scienceDocs, err := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestScienceCollection, + }) + assert.NoError(t, err) + t.Logf("✓ Science collection: %s (%d documents)", SearchTestScienceCollection, len(scienceDocs.Data)) + for _, doc := range scienceDocs.Data { + t.Logf(" - %s", doc["name"]) + } + + // Check Tech collection + techExists, err := kb.API.CollectionExists(ctx, SearchTestTechCollection) + if err != nil { + t.Fatalf("Failed to check tech collection: %v", err) + } + if !techExists.Exists { + t.Fatalf("✗ Tech collection does not exist. Run TestSearchSetup first.") + } + + techDocs, err := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: SearchTestTechCollection, + }) + assert.NoError(t, err) + t.Logf("✓ Tech collection: %s (%d documents)", SearchTestTechCollection, len(techDocs.Data)) + for _, doc := range techDocs.Data { + t.Logf(" - %s", doc["name"]) + } + + t.Log("") + t.Log("✓ Test data verified! Ready for search tests.") +} + +// ========== Helper Functions ========== + +// addFixedTestDocument adds a document for search testing +func addFixedTestDocument(t *testing.T, ctx context.Context, collectionID, title, content string) string { + params := &api.AddTextParams{ + CollectionID: collectionID, + Text: content, + DocID: fmt.Sprintf("%s__%s", collectionID, sanitizeTitle(title)), + Metadata: map[string]interface{}{ + "title": title, + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + // Enable extraction for graph-based search + Extraction: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "gpt-4o-mini", + }, + } + + result, err := kb.API.AddText(ctx, params) + if err != nil { + t.Logf("Warning: Failed to add document '%s': %v", title, err) + return "" + } + return result.DocID +} + +// sanitizeTitle converts title to a safe ID format +func sanitizeTitle(title string) string { + result := "" + for _, c := range title { + if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') { + result += string(c) + } else if c == ' ' { + result += "_" + } + } + return result +} diff --git a/kb/api/search_test.go b/kb/api/search_test.go new file mode 100644 index 00000000..646b9804 --- /dev/null +++ b/kb/api/search_test.go @@ -0,0 +1,446 @@ +package api_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + graphragtypes "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// Note: TestMain is defined in collection_test.go +// Note: Test data setup is in search_setup_test.go + +// ========== Search Query Tests ========== +// These tests use fixed collection IDs from search_setup_test.go +// Run TestSearchSetup first to create test data, then run these tests. +// +// Usage: +// 1. Setup (run once): go test -v -run "TestSearchSetup" ./kb/api/... +// 2. Run tests: go test -v -run "TestSearchQuery" ./kb/api/... +// 3. Cleanup (optional): go test -v -run "TestSearchCleanup" ./kb/api/... + +// verifyTestDataExists checks if test collections exist, skips if not +func verifyTestDataExists(t *testing.T, ctx context.Context) { + scienceExists, _ := kb.API.CollectionExists(ctx, SearchTestScienceCollection) + techExists, _ := kb.API.CollectionExists(ctx, SearchTestTechCollection) + + if scienceExists == nil || !scienceExists.Exists || techExists == nil || !techExists.Exists { + t.Skip("Test data not found. Run 'go test -v -run TestSearchSetup ./kb/api/...' first") + } +} + +func TestSearchQuery(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + verifyTestDataExists(t, ctx) + + t.Run("VectorSearch_SingleCollection", func(t *testing.T) { + // Test: Simple vector search in science collection + // Query about Einstein should find Einstein-related documents + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "Who is Albert Einstein and what did he discover?", + Mode: api.SearchModeVector, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Search error (may be expected if not implemented): %v", err) + return + } + + if result == nil { + t.Skip("Search not implemented yet (returned nil)") + } + + assert.Greater(t, len(result.Segments), 0, "Should find segments about Einstein") + t.Logf("Vector search returned %d segments", len(result.Segments)) + + // Verify relevance - top results should mention Einstein + for i, seg := range result.Segments { + t.Logf(" Segment %d (score: %.4f): %s...", i, seg.Score, truncateText(seg.Text, 100)) + } + }) + + t.Run("VectorSearch_MultipleQueries", func(t *testing.T) { + // Test: Multiple queries in same collection, results should be merged + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "relativity theory", + Mode: api.SearchModeVector, + PageSize: 3, + }, + { + CollectionID: SearchTestScienceCollection, + Input: "Nobel Prize physics", + Mode: api.SearchModeVector, + PageSize: 3, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Multi-query search returned %d merged segments", len(result.Segments)) + }) + + t.Run("VectorSearch_CrossCollection", func(t *testing.T) { + // Test: Search across both collections + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "innovation and discovery", + Mode: api.SearchModeVector, + PageSize: 3, + }, + { + CollectionID: SearchTestTechCollection, + Input: "technology innovation", + Mode: api.SearchModeVector, + PageSize: 3, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Cross-collection search returned %d segments", len(result.Segments)) + }) + + t.Run("ExpandSearch_EntityExpansion", func(t *testing.T) { + // Test: Expand mode should find related entities through graph + // Query: "photoelectric effect" should expand to find: + // - Einstein (discovered it) + // - Nobel Prize (awarded for it) + // - Quantum mechanics (built upon it) + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "photoelectric effect", + Mode: api.SearchModeExpand, + MaxDepth: 2, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Expand search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Expand search returned %d segments", len(result.Segments)) + + // Check if graph data is returned + if result.Graph != nil { + t.Logf(" Graph nodes: %d, relationships: %d", + len(result.Graph.Nodes), len(result.Graph.Relationships)) + } + + // Verify expanded results include related entities + for i, seg := range result.Segments { + t.Logf(" Segment %d (score: %.4f): %s...", i, seg.Score, truncateText(seg.Text, 100)) + } + }) + + t.Run("ExpandSearch_DeepAssociation", func(t *testing.T) { + // Test: Deep association through entity relationships + // Query: "Germany physics" should expand to find: + // - Einstein (born in Germany, physicist) + // - Relativity (Einstein's theory) + // - Planck (German physicist, quantum theory) + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "German physicist contributions", + Mode: api.SearchModeExpand, + MaxDepth: 3, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Deep expand search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Deep expand search returned %d segments", len(result.Segments)) + }) + + t.Run("GraphSearch_EntityTraversal", func(t *testing.T) { + // Test: Pure graph search - find segments through entity relationships + queries := []api.Query{ + { + CollectionID: SearchTestTechCollection, + Input: "Steve Jobs", + Mode: api.SearchModeGraph, + MaxDepth: 2, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Graph search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Graph search returned %d segments", len(result.Segments)) + + if result.Graph != nil { + t.Logf(" Found %d nodes, %d relationships", + len(result.Graph.Nodes), len(result.Graph.Relationships)) + for _, node := range result.Graph.Nodes { + t.Logf(" Node: %s (%s)", node.ID, node.EntityType) + } + } + }) + + t.Run("Search_WithMessages", func(t *testing.T) { + // Test: Search using conversation history instead of direct input + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Messages: []graphragtypes.ChatMessage{ + {Role: "user", Content: "Tell me about famous physicists"}, + {Role: "assistant", Content: "There are many famous physicists throughout history..."}, + {Role: "user", Content: "What about Einstein specifically?"}, + }, + Mode: api.SearchModeVector, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Message-based search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Message-based search returned %d segments", len(result.Segments)) + }) + + t.Run("Search_WithDocumentFilter", func(t *testing.T) { + // Test: Search within a specific document + // First, get a document ID + filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 1, + CollectionID: SearchTestScienceCollection, + } + listResult, err := kb.API.ListDocuments(ctx, filter) + if err != nil || len(listResult.Data) == 0 { + t.Skip("No documents available for filter test") + } + + docID, ok := listResult.Data[0]["document_id"].(string) + if !ok { + t.Skip("Could not get document ID") + } + + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + DocumentID: docID, + Input: "physics discovery", + Mode: api.SearchModeVector, + PageSize: 5, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Document-filtered search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Document-filtered search returned %d segments", len(result.Segments)) + + // Verify all results are from the specified document + for _, seg := range result.Segments { + if seg.DocumentID != "" { + assert.Equal(t, docID, seg.DocumentID, "All segments should be from filtered document") + } + } + }) + + t.Run("Search_WithPagination", func(t *testing.T) { + // Test: Pagination + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "physics", + Mode: api.SearchModeVector, + Page: 1, + PageSize: 2, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Paginated search error: %v", err) + return + } + + assert.NotNil(t, result) + assert.LessOrEqual(t, len(result.Segments), 2, "Should respect page size") + t.Logf("Page 1: %d segments, Total: %d, TotalPages: %d", + len(result.Segments), result.Total, result.TotalPages) + + // Get page 2 + queries[0].Page = 2 + result2, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Page 2 search error: %v", err) + return + } + + if result2 != nil && len(result2.Segments) > 0 { + t.Logf("Page 2: %d segments", len(result2.Segments)) + } + }) + + t.Run("Search_WithMinScore", func(t *testing.T) { + // Test: Filter by minimum score + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "Einstein relativity", + Mode: api.SearchModeVector, + MinScore: 0.5, + PageSize: 10, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("MinScore search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("MinScore search returned %d segments", len(result.Segments)) + + // Verify all results meet minimum score + for _, seg := range result.Segments { + assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet minimum score") + } + }) + + t.Run("Search_WithMetadataFilter", func(t *testing.T) { + // Test: Filter by metadata + queries := []api.Query{ + { + CollectionID: SearchTestScienceCollection, + Input: "physics", + Mode: api.SearchModeVector, + Metadata: map[string]interface{}{ + "title": "Albert Einstein Biography", + }, + PageSize: 10, + }, + } + + result, err := kb.API.Search(ctx, queries) + if err != nil { + t.Logf("Metadata filter search error: %v", err) + return + } + + assert.NotNil(t, result) + t.Logf("Metadata-filtered search returned %d segments", len(result.Segments)) + }) +} + +// ========== Error Handling Tests ========== + +func TestSearchErrorHandling(t *testing.T) { + if kb.API == nil { + t.Skip("KB API not initialized") + } + + ctx := context.Background() + + t.Run("EmptyQueries", func(t *testing.T) { + result, err := kb.API.Search(ctx, []api.Query{}) + // Empty queries should return empty result or error + if err != nil { + assert.Contains(t, err.Error(), "required") + } else { + assert.NotNil(t, result) + assert.Equal(t, 0, len(result.Segments)) + } + }) + + t.Run("MissingCollectionID", func(t *testing.T) { + queries := []api.Query{ + { + Input: "test query", + Mode: api.SearchModeVector, + }, + } + + _, err := kb.API.Search(ctx, queries) + assert.Error(t, err) + assert.Contains(t, err.Error(), "collection") + }) + + t.Run("MissingInputAndMessages", func(t *testing.T) { + queries := []api.Query{ + { + CollectionID: "some_collection", + Mode: api.SearchModeVector, + }, + } + + _, err := kb.API.Search(ctx, queries) + assert.Error(t, err) + assert.Contains(t, err.Error(), "input") + }) + + t.Run("NonexistentCollection", func(t *testing.T) { + queries := []api.Query{ + { + CollectionID: "nonexistent_collection_xyz", + Input: "test query", + Mode: api.SearchModeVector, + }, + } + + _, err := kb.API.Search(ctx, queries) + assert.Error(t, err) + }) +} + +// ========== Helper Functions ========== + +func truncateText(text string, maxLen int) string { + if len(text) <= maxLen { + return text + } + return text[:maxLen] + "..." +} diff --git a/kb/api/types.go b/kb/api/types.go index a2997b26..05e9e95e 100644 --- a/kb/api/types.go +++ b/kb/api/types.go @@ -195,3 +195,80 @@ type AddDocumentAsyncResult struct { JobID string `json:"job_id" yaml:"job_id"` DocID string `json:"doc_id" yaml:"doc_id"` } + +// ========== Search Types ========== + +// SearchMode defines the search strategy +type SearchMode string + +const ( + // SearchModeVector performs pure vector similarity search + SearchModeVector SearchMode = "vector" + // SearchModeGraph performs graph traversal to find related segments + SearchModeGraph SearchMode = "graph" + // SearchModeExpand uses graph to expand/associate entities, then enhances vector search + // This enables deeper semantic connections through entity relationships + SearchModeExpand SearchMode = "expand" +) + +// Query represents a single search query +type Query struct { + // CollectionID is the collection to search in (required) + CollectionID string `json:"collection_id" yaml:"collection_id"` + + // Input is the direct search query text (e.g., LLM-summarized query) + // Either Input or Messages is required; Input takes precedence if both provided + Input string `json:"input,omitempty" yaml:"input,omitempty"` + + // Messages is the conversation history for context-aware search + // The last user message is used as the query if Input is empty + Messages []types.ChatMessage `json:"messages,omitempty" yaml:"messages,omitempty"` + + // Mode determines the search strategy (optional, defaults to collection config or "expand") + // - vector: pure vector similarity search + // - graph: graph traversal to find related segments + // - expand: graph-based entity expansion/association + vector search + Mode SearchMode `json:"mode,omitempty" yaml:"mode,omitempty"` + + // DocumentID filters results to a specific document (optional) + DocumentID string `json:"document_id,omitempty" yaml:"document_id,omitempty"` + + // MinScore filters results below this similarity threshold (optional) + MinScore float64 `json:"min_score,omitempty" yaml:"min_score,omitempty"` + + // Metadata filters segments by metadata fields (optional) + Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` + + // Graph search options (used when Mode is graph or hybrid) + MaxDepth int `json:"max_depth,omitempty" yaml:"max_depth,omitempty"` // Max traversal depth (default: 2) + + // Pagination options + // If not specified, returns default number of results + Page int `json:"page,omitempty" yaml:"page,omitempty"` // Page number (1-based), 0 means no pagination + PageSize int `json:"pagesize,omitempty" yaml:"pagesize,omitempty"` // Number of results per page + Cursor string `json:"cursor,omitempty" yaml:"cursor,omitempty"` // Cursor for cursor-based pagination +} + +// GraphData contains graph-specific search results +type GraphData struct { + Nodes []types.GraphNode `json:"nodes,omitempty" yaml:"nodes,omitempty"` + Relationships []types.GraphRelationship `json:"relationships,omitempty" yaml:"relationships,omitempty"` +} + +// SearchResult represents the merged result of search operations +type SearchResult struct { + // Segments contains the merged and deduplicated text segments with scores + Segments []types.Segment `json:"segments" yaml:"segments"` + + // Graph contains merged nodes and relationships (only for graph/hybrid mode) + Graph *GraphData `json:"graph,omitempty" yaml:"graph,omitempty"` + + // Pagination info + Page int `json:"page,omitempty" yaml:"page,omitempty"` // Current page number + PageSize int `json:"pagesize,omitempty" yaml:"pagesize,omitempty"` // Results per page + Total int `json:"total" yaml:"total"` // Total number of results + TotalPages int `json:"pagecnt,omitempty" yaml:"pagecnt,omitempty"` // Total pages + Next int `json:"next,omitempty" yaml:"next,omitempty"` // Next page number + Prev int `json:"prev,omitempty" yaml:"prev,omitempty"` // Previous page number + Cursor string `json:"cursor,omitempty" yaml:"cursor,omitempty"` // Cursor for next page +} From 09d368dae7a45f6de8c9f9ec6c54b1d0e5384560 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 20 Dec 2025 11:44:09 +0800 Subject: [PATCH 08/10] Add KB Tests and Update Makefile for Test Management - Introduced a new unit test target for KB tests in the Makefile, allowing for dedicated testing of the KB module. - Updated the test folder selection logic to exclude additional AI-related components, ensuring focused testing. - Enhanced the GitHub Actions workflows to include KB tests, setting up necessary services like Qdrant, Neo4j, and MongoDB for a comprehensive testing environment. - Refactored search test functions to improve data existence checks and streamline test setup processes, enhancing test reliability and maintainability. --- .github/workflows/pr-test.yml | 186 ++++++++++++++++++ .github/workflows/unit-test.yml | 124 ++++++++++++ Makefile | 38 +++- kb/api/README.md | 334 ++++++++++++++++++++++++++++++++ kb/api/search.go | 13 +- kb/api/search_test.go | 22 +-- kb/api/types.go | 2 +- 7 files changed, 696 insertions(+), 23 deletions(-) create mode 100644 kb/api/README.md diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 1c541061..e12c9b39 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -164,6 +164,192 @@ env: TWILIO_TEST_PHONE: ${{ secrets.TWILIO_TEST_PHONE }} jobs: + # ============================================================================= + # KB Tests (kb) - Run once with SQLite (requires Qdrant, Neo4j, FastEmbed) + # ============================================================================= + KBTest: + runs-on: ubuntu-latest + services: + qdrant: + image: qdrant/qdrant:latest + ports: + - 6333:6333 + - 6334:6334 + + fastembed: + image: yaoapp/fastembed:latest-amd64 + env: + FASTEMBED_PASSWORD: Yao@2026 + ports: + - 6001:8000 + + neo4j: + image: neo4j:latest + ports: + - "7687:7687" + env: + NEO4J_AUTH: neo4j/Yao2026Neo4j + + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + + strategy: + matrix: + go: [1.24] + if: > + ${{ github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.conclusion == 'success' }} + steps: + - name: "Download artifact" + uses: actions/github-script@v7 + with: + script: | + var artifacts = await github.rest.actions.listWorkflowRunArtifacts({ + owner: context.repo.owner, + repo: context.repo.repo, + run_id: ${{github.event.workflow_run.id }}, + }); + var matchArtifact = artifacts.data.artifacts.filter((artifact) => { + return artifact.name == "pr" + })[0]; + var download = await github.rest.actions.downloadArtifact({ + owner: context.repo.owner, + repo: context.repo.repo, + artifact_id: matchArtifact.id, + archive_format: 'zip', + }); + var fs = require('fs'); + fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data)); + + - name: "Read NR & SHA" + run: | + unzip pr.zip + cat NR + cat SHA + echo HEAD=$(cat SHA) >> $GITHUB_ENV + echo NR=$(cat NR) >> $GITHUB_ENV + + - name: "Comment on PR" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '🤖 KB Tests (kb) running with SQLite...' + }); + + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: yaoapp/kun + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: yaoapp/xun + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: yaoapp/gou + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout pull request HEAD commit + uses: actions/checkout@v4 + with: + ref: ${{ env.HEAD }} + + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Start Redis + uses: supercharge/redis-github-action@1.4.0 + with: + redis-version: 6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + + - name: Run KB Tests (kb) + run: make unit-test-kb + + - name: "Comment on PR - KB Tests Done" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '✅ KB Tests (kb) passed!' + }); + # ============================================================================= # AI Tests (agent, aigc) - Run once with SQLite # ============================================================================= diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index c3bf659d..297ddfb6 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -171,6 +171,130 @@ env: TWILIO_TEST_PHONE: ${{ secrets.TWILIO_TEST_PHONE }} jobs: + # ============================================================================= + # KB Tests (kb) - Run once with SQLite (requires Qdrant, Neo4j, FastEmbed) + # ============================================================================= + kb-test: + runs-on: ubuntu-latest + services: + qdrant: + image: qdrant/qdrant:latest + ports: + - 6333:6333 + - 6334:6334 + + fastembed: + image: yaoapp/fastembed:latest-amd64 + env: + FASTEMBED_PASSWORD: Yao@2026 + ports: + - 6001:8000 + + neo4j: + image: neo4j:latest + ports: + - "7687:7687" + env: + NEO4J_AUTH: neo4j/Yao2026Neo4j + + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + + strategy: + matrix: + go: [1.24] + steps: + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_KUN }} + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_XUN }} + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_GOU }} + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout Code + uses: actions/checkout@v4 + + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Start Redis + uses: supercharge/redis-github-action@1.4.0 + with: + redis-version: 6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + + - name: Run KB Tests (kb) + run: make unit-test-kb + # ============================================================================= # AI Tests (agent, aigc) - Run once with SQLite # ============================================================================= diff --git a/Makefile b/Makefile index 5e36c619..785b1c67 100644 --- a/Makefile +++ b/Makefile @@ -11,10 +11,12 @@ OS := $(shell uname) # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*' | awk '!/\/tests\// || /openapi\/tests/') -# Core tests (exclude AI-related: agent, aigc, openai) -TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent' | awk '!/\/tests\// || /openapi\/tests/') +# Core tests (exclude AI-related: agent, aigc, openai, and KB) +TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb' | awk '!/\/tests\// || /openapi\/tests/') # AI tests (agent, aigc) TESTFOLDER_AI := $(shell $(GO) list ./agent/... ./aigc/...) +# KB tests (kb) +TESTFOLDER_KB := $(shell $(GO) list ./kb/...) TESTTAGS ?= "" # TESTWIDGETS := $(shell $(GO) list ./widgets/...) @@ -103,6 +105,38 @@ unit-test-ai: fi; \ done +# KB Unit Test (kb) +.PHONY: unit-test-kb +unit-test-kb: + echo "mode: count" > coverage-kb.out + for d in $(TESTFOLDER_KB); do \ + $(GO) test -tags $(TESTTAGS) -v -timeout=20m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestSearchCleanup' $$d > tmp.out; \ + cat tmp.out; \ + if grep -q "^--- FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "^FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "^panic:" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "build failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "setup failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "runtime error" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + fi; \ + if [ -f profile.out ]; then \ + cat profile.out | grep -v "mode:" >> coverage-kb.out; \ + rm profile.out; \ + fi; \ + done + # Benchmark Test .PHONY: benchmark benchmark: diff --git a/kb/api/README.md b/kb/api/README.md new file mode 100644 index 00000000..3742ddf0 --- /dev/null +++ b/kb/api/README.md @@ -0,0 +1,334 @@ +# KB API + +The `kb/api` package provides a unified Go API for Knowledge Base operations including collection management, document ingestion, and semantic search. + +## Quick Start + +```go +import ( + "context" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/kb/api" +) + +// After kb.Load(), use kb.API to access all operations +ctx := context.Background() +``` + +## API Interface + +```go +type API interface { + // Collection operations + CreateCollection(ctx, params) (*CreateCollectionResult, error) + RemoveCollection(ctx, collectionID) (*RemoveCollectionResult, error) + GetCollection(ctx, collectionID) (map[string]interface{}, error) + CollectionExists(ctx, collectionID) (*CollectionExistsResult, error) + ListCollections(ctx, filter) (*ListCollectionsResult, error) + UpdateCollectionMetadata(ctx, collectionID, params) (*UpdateMetadataResult, error) + + // Document operations + ListDocuments(ctx, filter) (*ListDocumentsResult, error) + GetDocument(ctx, docID, params) (map[string]interface{}, error) + RemoveDocuments(ctx, params) (*RemoveDocumentsResult, error) + + // Document add operations (sync) + AddFile(ctx, params) (*AddDocumentResult, error) + AddText(ctx, params) (*AddDocumentResult, error) + AddURL(ctx, params) (*AddDocumentResult, error) + + // Document add operations (async) + AddFileAsync(ctx, params) (*AddDocumentAsyncResult, error) + AddTextAsync(ctx, params) (*AddDocumentAsyncResult, error) + AddURLAsync(ctx, params) (*AddDocumentAsyncResult, error) + + // Search operations + Search(ctx, queries) (*SearchResult, error) +} +``` + +## Collection Operations + +### Create Collection + +```go +params := &api.CreateCollectionParams{ + ID: "my_collection", + Metadata: map[string]interface{}{ + "name": "My Knowledge Base", + "description": "Collection description", + }, + EmbeddingProviderID: "__yao.openai", + EmbeddingOptionID: "text-embedding-3-small", + Locale: "en", + Config: &types.CreateCollectionOptions{ + Distance: "cosine", + IndexType: "hnsw", + }, +} + +result, err := kb.API.CreateCollection(ctx, params) +// result.CollectionID = "my_collection" +``` + +### Get Collection + +```go +collection, err := kb.API.GetCollection(ctx, "my_collection") +// collection["id"], collection["name"], collection["config"], etc. +``` + +### List Collections + +```go +filter := &api.ListCollectionsFilter{ + Page: 1, + PageSize: 20, + Keywords: "knowledge", + Status: []string{"active"}, +} + +result, err := kb.API.ListCollections(ctx, filter) +// result.Data, result.Total, result.PageCnt +``` + +### Remove Collection + +```go +result, err := kb.API.RemoveCollection(ctx, "my_collection") +// result.Removed = true +``` + +## Document Operations + +### Add Text + +```go +params := &api.AddTextParams{ + CollectionID: "my_collection", + Text: "Einstein developed the theory of relativity...", + DocID: "einstein_bio", // optional, auto-generated if empty + Metadata: map[string]interface{}{ + "title": "Einstein Biography", + "author": "John Doe", + }, + Chunking: &api.ProviderConfigParams{ + ProviderID: "__yao.structured", + OptionID: "standard", + }, + Embedding: &api.ProviderConfigParams{ + ProviderID: "__yao.openai", + OptionID: "text-embedding-3-small", + }, + Extraction: &api.ProviderConfigParams{ // optional, for graph extraction + ProviderID: "__yao.openai", + OptionID: "gpt-4o-mini", + }, +} + +result, err := kb.API.AddText(ctx, params) +// result.DocID = "einstein_bio" +``` + +### Add File + +```go +params := &api.AddFileParams{ + CollectionID: "my_collection", + FileID: "uploaded_file_id", + Uploader: "local", // or "s3", etc. + Chunking: &api.ProviderConfigParams{...}, + Embedding: &api.ProviderConfigParams{...}, +} + +result, err := kb.API.AddFile(ctx, params) +``` + +### Add URL + +```go +params := &api.AddURLParams{ + CollectionID: "my_collection", + URL: "https://example.com/article", + Chunking: &api.ProviderConfigParams{...}, + Embedding: &api.ProviderConfigParams{...}, +} + +result, err := kb.API.AddURL(ctx, params) +``` + +### List Documents + +```go +filter := &api.ListDocumentsFilter{ + Page: 1, + PageSize: 20, + CollectionID: "my_collection", + Status: []string{"active"}, +} + +result, err := kb.API.ListDocuments(ctx, filter) +``` + +### Remove Documents + +```go +params := &api.RemoveDocumentsParams{ + DocumentIDs: []string{"doc1", "doc2"}, +} + +result, err := kb.API.RemoveDocuments(ctx, params) +``` + +## Search Operations + +The Search API supports batch queries with three search modes: + +| Mode | Description | +|------|-------------| +| `vector` | Pure vector similarity search | +| `graph` | Graph traversal to find related segments via entities | +| `expand` | Graph-based entity expansion + vector search (default) | + +### Basic Vector Search + +```go +queries := []api.Query{ + { + CollectionID: "my_collection", + Input: "What is the theory of relativity?", + Mode: api.SearchModeVector, + PageSize: 10, + }, +} + +result, err := kb.API.Search(ctx, queries) +// result.Segments - matched text segments with scores +// result.Total - total count +``` + +### Graph-Enhanced Search (Expand Mode) + +```go +queries := []api.Query{ + { + CollectionID: "my_collection", + Input: "Einstein's contributions to physics", + Mode: api.SearchModeExpand, // default + MaxDepth: 2, // graph traversal depth + PageSize: 10, + }, +} + +result, err := kb.API.Search(ctx, queries) +// result.Segments - segments from vector + graph expansion +// result.Graph.Nodes - related entities +// result.Graph.Relationships - entity relationships +``` + +### Multi-Query Search + +Queries can span multiple collections; results are merged and deduplicated: + +```go +queries := []api.Query{ + { + CollectionID: "science_kb", + Input: "quantum mechanics", + Mode: api.SearchModeVector, + }, + { + CollectionID: "tech_kb", + Input: "machine learning", + Mode: api.SearchModeVector, + }, +} + +result, err := kb.API.Search(ctx, queries) +// Merged results from both collections +``` + +### Search with Messages (Conversation Context) + +```go +queries := []api.Query{ + { + CollectionID: "my_collection", + Messages: []types.ChatMessage{ + {Role: "user", Content: "Tell me about Einstein"}, + {Role: "assistant", Content: "Einstein was a physicist..."}, + {Role: "user", Content: "What about his discoveries?"}, // used as query + }, + Mode: api.SearchModeExpand, + }, +} + +result, err := kb.API.Search(ctx, queries) +``` + +### Search with Filters + +```go +queries := []api.Query{ + { + CollectionID: "my_collection", + Input: "physics", + DocumentID: "specific_doc_id", // filter to specific document + MinScore: 0.5, // minimum similarity score + Metadata: map[string]interface{}{ + "category": "science", + }, + Page: 1, + PageSize: 20, + }, +} + +result, err := kb.API.Search(ctx, queries) +``` + +## Query Parameters + +| Field | Type | Description | +|-------|------|-------------| +| `CollectionID` | string | Collection to search (required) | +| `Input` | string | Direct query text | +| `Messages` | []ChatMessage | Conversation history (last user message used as query) | +| `Mode` | SearchMode | `vector`, `graph`, or `expand` (default: `expand`) | +| `DocumentID` | string | Filter to specific document | +| `MinScore` | float64 | Minimum similarity threshold | +| `Metadata` | map | Filter by metadata fields | +| `MaxDepth` | int | Graph traversal depth (default: 2) | +| `Page` | int | Page number (1-based) | +| `PageSize` | int | Results per page | + +## Search Result + +```go +type SearchResult struct { + Segments []types.Segment // Matched segments with scores + Graph *GraphData // Nodes and relationships (graph/expand mode) + Total int // Total results count + Page int // Current page + PageSize int // Results per page + TotalPages int // Total pages + Next int // Next page number + Prev int // Previous page number +} +``` + +## Provider Configuration + +Providers handle text processing: + +```go +type ProviderConfigParams struct { + ProviderID string // e.g., "__yao.openai", "__yao.structured" + OptionID string // e.g., "text-embedding-3-small", "gpt-4o-mini" +} +``` + +Common providers: +- **Chunking**: `__yao.structured` - text splitting +- **Embedding**: `__yao.openai` - vector embeddings +- **Extraction**: `__yao.openai` - entity/relationship extraction for graph + diff --git a/kb/api/search.go b/kb/api/search.go index cc1bfe4d..f1cb3793 100644 --- a/kb/api/search.go +++ b/kb/api/search.go @@ -73,12 +73,17 @@ func (kb *KBInstance) Search(ctx context.Context, queries []Query) (*SearchResul close(errChan) // Collect errors - var errors []error + var errs []error for err := range errChan { - errors = append(errors, err) + errs = append(errs, err) } - if len(errors) > 0 { - log.Warn("Search completed with errors: %v", errors) + if len(errs) > 0 { + // Combine all error messages + errMsgs := make([]string, len(errs)) + for i, e := range errs { + errMsgs[i] = e.Error() + } + return nil, fmt.Errorf("search failed: %v", errMsgs) } // 4. Merge and deduplicate results diff --git a/kb/api/search_test.go b/kb/api/search_test.go index 646b9804..91284d65 100644 --- a/kb/api/search_test.go +++ b/kb/api/search_test.go @@ -14,22 +14,12 @@ import ( // Note: Test data setup is in search_setup_test.go // ========== Search Query Tests ========== -// These tests use fixed collection IDs from search_setup_test.go -// Run TestSearchSetup first to create test data, then run these tests. -// -// Usage: -// 1. Setup (run once): go test -v -run "TestSearchSetup" ./kb/api/... -// 2. Run tests: go test -v -run "TestSearchQuery" ./kb/api/... -// 3. Cleanup (optional): go test -v -run "TestSearchCleanup" ./kb/api/... -// verifyTestDataExists checks if test collections exist, skips if not -func verifyTestDataExists(t *testing.T, ctx context.Context) { - scienceExists, _ := kb.API.CollectionExists(ctx, SearchTestScienceCollection) - techExists, _ := kb.API.CollectionExists(ctx, SearchTestTechCollection) - - if scienceExists == nil || !scienceExists.Exists || techExists == nil || !techExists.Exists { - t.Skip("Test data not found. Run 'go test -v -run TestSearchSetup ./kb/api/...' first") - } +// ensureTestDataExists ensures test collections exist by running setup if needed +// Setup will skip creation if data already exists +func ensureTestDataExists(t *testing.T, ctx context.Context) { + // Run setup - it checks if data exists and skips if already complete + TestSearchSetup(t) } func TestSearchQuery(t *testing.T) { @@ -38,7 +28,7 @@ func TestSearchQuery(t *testing.T) { } ctx := context.Background() - verifyTestDataExists(t, ctx) + ensureTestDataExists(t, ctx) t.Run("VectorSearch_SingleCollection", func(t *testing.T) { // Test: Simple vector search in science collection diff --git a/kb/api/types.go b/kb/api/types.go index 05e9e95e..6ae047eb 100644 --- a/kb/api/types.go +++ b/kb/api/types.go @@ -239,7 +239,7 @@ type Query struct { // Metadata filters segments by metadata fields (optional) Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` - // Graph search options (used when Mode is graph or hybrid) + // Graph search options (used when Mode is graph or expand) MaxDepth int `json:"max_depth,omitempty" yaml:"max_depth,omitempty"` // Max traversal depth (default: 2) // Pagination options From f8bcea23fce7d26d4e3707c664f1f754c5250776 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 20 Dec 2025 11:54:53 +0800 Subject: [PATCH 09/10] Update GitHub Actions workflows to use Go version 1.25 - Changed the Go version in multiple workflow files from 1.24 to 1.25, ensuring compatibility with the latest features and improvements in the Go programming language. --- .github/workflows/pr-test.yml | 8 ++++---- .github/workflows/unit-test.yml | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index e12c9b39..cf32f336 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -201,7 +201,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] if: > ${{ github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.conclusion == 'success' }} @@ -393,7 +393,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] if: > ${{ github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.conclusion == 'success' }} @@ -565,7 +565,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] db: [MySQL8.0, SQLite3] if: > ${{ github.event.workflow_run.event == 'pull_request' && @@ -733,7 +733,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] db: [MySQL8.0, SQLite3] redis: [4, 5, 6] mongo: ["6.0"] diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 297ddfb6..2935e50b 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -208,7 +208,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] steps: - name: Checkout Kun uses: actions/checkout@v4 @@ -338,7 +338,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] steps: - name: Checkout Kun uses: actions/checkout@v4 @@ -448,7 +448,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] db: [MySQL8.0, SQLite3] steps: - name: Checkout Kun @@ -581,7 +581,7 @@ jobs: strategy: matrix: - go: [1.24] + go: ["1.25"] db: [MySQL8.0, SQLite3] redis: [4, 5, 6] mongo: ["6.0"] From f58aea74590317f5b132f539259132e694f28396 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 20 Dec 2025 12:12:57 +0800 Subject: [PATCH 10/10] Update Dockerfiles and GitHub Actions to use version 1.0.0 - Updated the Dockerfiles and GitHub Actions workflows to reference version 1.0.0 of the Yao build environment and Go, ensuring compatibility with the latest features and improvements. - Adjusted the Go version in the macOS workflow and updated the Dockerfile to install Go 1.25.0, enhancing the build process. --- .github/workflows/build-linux.yml | 2 +- .github/workflows/build-macos.yml | 4 ++-- docker/build/Dockerfile | 16 ++++++++-------- docker/development/Dockerfile | 8 ++++---- docker/production/Dockerfile | 10 +++++----- 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/.github/workflows/build-linux.yml b/.github/workflows/build-linux.yml index d3371aed..2a2ad563 100644 --- a/.github/workflows/build-linux.yml +++ b/.github/workflows/build-linux.yml @@ -10,7 +10,7 @@ jobs: build: runs-on: "ubuntu-latest" container: - image: yaoapp/yao-build:0.10.5 + image: yaoapp/yao-build:1.0.0 env: CF_ACCESS_KEY_ID: ${{ secrets.CF_ACCESS_KEY_ID }} diff --git a/.github/workflows/build-macos.yml b/.github/workflows/build-macos.yml index 4fff33de..abba7b72 100644 --- a/.github/workflows/build-macos.yml +++ b/.github/workflows/build-macos.yml @@ -7,13 +7,13 @@ on: description: "Version tags" env: - VERSION: 0.10.5 + VERSION: 1.0.0 jobs: build: strategy: matrix: - go: [1.24.3] + go: ["1.25"] runs-on: "macos-latest" steps: - name: Setup Node.js diff --git a/docker/build/Dockerfile b/docker/build/Dockerfile index 9e1ff9e9..92ae1159 100644 --- a/docker/build/Dockerfile +++ b/docker/build/Dockerfile @@ -2,14 +2,14 @@ # Yao Build Environment (Ubuntu 24.04 AMD64) # # Build: -# docker build --platform linux/amd64 -t yaoapp/yao-build:0.10.5 . +# docker build --platform linux/amd64 -t yaoapp/yao-build:1.0.0 . # # Usage: -# docker run --rm -it -v /local/path/dist:/data yaoapp/yao-build:0.10.5 +# docker run --rm -it -v /local/path/dist:/data yaoapp/yao-build:1.0.0 # # Tests: -# docker run --rm -it yaoapp/yao-build:0.10.5 /bin/bash -# docker run --rm -it -v ./test:/data yaoapp/yao-build:0.10.5 /bin/bash +# docker run --rm -it yaoapp/yao-build:1.0.0 /bin/bash +# docker run --rm -it -v ./test:/data yaoapp/yao-build:1.0.0 /bin/bash # # =========================================== FROM ubuntu:24.04 @@ -29,10 +29,10 @@ RUN apt-get update && \ apt-get install -y git && \ apt-get install -y unzip -# Install Go 1.24.3 -RUN wget https://golang.org/dl/go1.24.3.linux-amd64.tar.gz && \ - tar -C /usr/local -xzf go1.24.3.linux-amd64.tar.gz && \ - rm go1.24.3.linux-amd64.tar.gz +# Install Go 1.25.0 +RUN wget https://golang.org/dl/go1.25.0.linux-amd64.tar.gz && \ + tar -C /usr/local -xzf go1.25.0.linux-amd64.tar.gz && \ + rm go1.25.0.linux-amd64.tar.gz # Install Node.js 18.x RUN curl -fsSL https://deb.nodesource.com/setup_18.x | bash - && \ diff --git a/docker/development/Dockerfile b/docker/development/Dockerfile index e8ffc323..ef778e28 100644 --- a/docker/development/Dockerfile +++ b/docker/development/Dockerfile @@ -6,12 +6,12 @@ # -t yaoapp/yao-dev:${VERSION}-${ARCH} . # # Build: -# docker build --platform linux/amd64 --build-arg VERSION=0.9.1 --build-arg ARCH=amd64 -t yaoapp/yao:0.9.1-amd64-dev . -# docker build --platform linux/arm64 --build-arg VERSION=0.9.1 --build-arg ARCH=arm64 -t yaoapp/yao:0.9.1-arm64-dev . +# docker build --platform linux/amd64 --build-arg VERSION=1.0.0 --build-arg ARCH=amd64 -t yaoapp/yao:1.0.0-amd64-dev . +# docker build --platform linux/arm64 --build-arg VERSION=1.0.0 --build-arg ARCH=arm64 -t yaoapp/yao:1.0.0-arm64-dev . # # Tests: -# docker run --rm yaoapp/yao:0.9.1-amd64-dev yao version -# docker run -d -p 5099:5099 yaoapp/yao:0.9.1-amd64-dev +# docker run --rm yaoapp/yao:1.0.0-amd64-dev yao version +# docker run -d -p 5099:5099 yaoapp/yao:1.0.0-amd64-dev # # =========================================== FROM ubuntu:24.04 diff --git a/docker/production/Dockerfile b/docker/production/Dockerfile index 8e788a80..68985492 100644 --- a/docker/production/Dockerfile +++ b/docker/production/Dockerfile @@ -3,15 +3,15 @@ # docker build \ # --build-arg VERSION="${VERSION}" \ # --build-arg ARCH="${ARCH}" \ -# -t yaoapp/yao-dev:${VERSION}-${ARCH} . +# -t yaoapp/yao:${VERSION}-${ARCH} . # # Build: -# docker build --platform linux/amd64 --build-arg VERSION=0.9.1 --build-arg ARCH=amd64 -t yaoapp/yao:0.9.1-amd64 . -# docker build --platform linux/arm64 --build-arg VERSION=0.9.1 --build-arg ARCH=arm64 -t yaoapp/yao:0.9.1-arm64 . +# docker build --platform linux/amd64 --build-arg VERSION=1.0.0 --build-arg ARCH=amd64 -t yaoapp/yao:1.0.0-amd64 . +# docker build --platform linux/arm64 --build-arg VERSION=1.0.0 --build-arg ARCH=arm64 -t yaoapp/yao:1.0.0-arm64 . # # Tests: -# docker run --rm yaoapp/yao:0.9.1-amd64 yao version -# docker run -d -p 5099:5099 yaoapp/yao:0.9.1-amd64 +# docker run --rm yaoapp/yao:1.0.0-amd64 yao version +# docker run -d -p 5099:5099 yaoapp/yao:1.0.0-amd64 # # =========================================== FROM alpine:latest