package kb import ( "context" "fmt" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/utils" "github.com/yaoapp/kun/log" "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi/response" ) // AddFileProcess processes a file addition request with business logic only // This function is Gin-agnostic and can be used for both sync and async operations func AddFileProcess(ctx context.Context, req *AddFileRequest, jobID ...string) error { // Check if kb.Instance is available if kb.Instance == nil { return fmt.Errorf("knowledge base not initialized") } // Validate request if err := req.Validate(); err != nil { return err } // 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) } // DocID should be generated by the caller before calling this function if req.DocID == "" { return fmt.Errorf("document ID is required") } // 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), } // Add job_id if provided (for async operations) if len(jobID) > 0 && jobID[0] != "" { documentData["job_id"] = jobID[0] } // Add base request fields req.BaseUpsertRequest.AddBaseFields(documentData) // First create database record _, err = config.CreateDocument(maps.MapStrAny(documentData)) if err != nil { return fmt.Errorf("failed to save document metadata: %w", err) } // Convert request to UpsertOptions upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions(path, contentType) if err != nil { // Rollback: remove the database record if rollbackErr := config.RemoveDocument(req.DocID); rollbackErr != nil { log.Error("Failed to rollback document database record: %v", rollbackErr) } 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 } // addFileWithRequest processes a file addition with pre-parsed request using Gin context func addFileWithRequest(c *gin.Context, req *AddFileRequest) { // Use the business logic function err := AddFileProcess(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": "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 func AddFile(c *gin.Context) { var req AddFileRequest // Check if kb.Instance is available if !checkKBInstance(c) { return } // Parse and bind JSON request if err := c.ShouldBindJSON(&req); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid request format: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate request if err := req.Validate(); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Generate document ID if not provided if req.DocID == "" { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } // Process the request addFileWithRequest(c, &req) } // AddFileAsync adds file to a collection asynchronously func AddFileAsync(c *gin.Context) { var req AddFileRequest // Check if kb.Instance is available if !checkKBInstance(c) { return } // Parse and bind JSON request if err := c.ShouldBindJSON(&req); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid request format: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate request if err := req.Validate(); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate file and get path _, _, err := validateFileAndGetPath(c, &req) if err != nil { return } // Convert request to UpsertOptions (just for validation) _, err = getUpsertOptions(c, &req.BaseUpsertRequest) if err != nil { return } // Generate document ID if not provided if req.DocID == "" { req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } // Create and run job job := NewJob() jobID := job.Run(func() { err := AddFileProcess(context.Background(), &req, job.ID) if err != nil { log.Error("Async file processing failed: %v", err) } }) // Return job_id and doc_id response.RespondWithSuccess(c, response.StatusCreated, gin.H{ "job_id": jobID, "doc_id": req.DocID, }) }