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