From 7b2534bf3710708346d0dc4f6b09cc6b83a66e58 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 14 Aug 2025 08:52:43 +0800 Subject: [PATCH] Refactor file, text, and URL addition handlers for improved request processing and validation - Introduced separate processing functions (AddFileProcess, ProcessAddTextRequest) to encapsulate business logic for file and text additions, making them Gin-agnostic. - Updated AddFile, AddText, and AddURL functions to handle request parsing and validation more effectively. - Enhanced error handling and rollback mechanisms during document creation and upsert operations. - Implemented async processing capabilities for AddFileAsync, AddTextAsync, and AddURLAsync functions using pre-parsed request data. - Streamlined utility functions for preparing requests and validating input, improving code maintainability and readability. --- openapi/kb/addfile.go | 190 ++++++++++++++++++++++++++++++----------- openapi/kb/addtext.go | 145 +++++++++++++++++++++++++++++-- openapi/kb/addurl.go | 70 ++++++++++++--- openapi/kb/document.go | 33 ++----- openapi/kb/utils.go | 32 +++---- 5 files changed, 354 insertions(+), 116 deletions(-) diff --git a/openapi/kb/addfile.go b/openapi/kb/addfile.go index 56b08e8d..d4ec337c 100644 --- a/openapi/kb/addfile.go +++ b/openapi/kb/addfile.go @@ -1,84 +1,104 @@ 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" ) -// AddFile adds a file to a collection -func AddFile(c *gin.Context) { +// 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) error { // Check if kb.Instance is available - if !checkKBInstance(c) { - return + if kb.Instance == nil { + return fmt.Errorf("knowledge base not initialized") } - // Prepare request and database data - req, documentData, err := PrepareAddFile(c) + // 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 { - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: err.Error(), - } - response.RespondWithError(c, response.StatusBadRequest, errorResp) - return + 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) + } + + // Generate document ID if not provided + if req.DocID == "" { + req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) } // Get KB config config, err := kb.GetConfig() if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to get KB config: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return + 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_name": fileInfo.Filename, + "file_path": path, + "file_mime_type": contentType, + "size": int64(fileInfo.Bytes), + } + + // Add base request fields + req.BaseUpsertRequest.AddBaseFields(documentData) + // First create database record _, err = config.CreateDocument(maps.MapStrAny(documentData)) if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to save document metadata: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return + return fmt.Errorf("failed to save document metadata: %w", err) } // Convert request to UpsertOptions - path, contentType, err := validateFileAndGetPath(c, req) + upsertOptions, err := req.BaseUpsertRequest.ToUpsertOptions(path, contentType) if err != nil { // Rollback: remove the database record - if err := config.RemoveDocument(req.DocID); err != nil { - log.Error("Failed to rollback document database record: %v", err) + if rollbackErr := config.RemoveDocument(req.DocID); rollbackErr != nil { + log.Error("Failed to rollback document database record: %v", rollbackErr) } - return + return fmt.Errorf("failed to convert request to upsert options: %w", err) } - upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest, path, contentType) + // Perform upsert operation with file path + _, err = kb.Instance.AddFile(ctx, path, upsertOptions) if err != nil { - // Rollback: remove the database record - if err := config.RemoveDocument(req.DocID); err != nil { - log.Error("Failed to rollback document database record: %v", err) - } - return - } - - // Perform upsert operation with file ID - _, err = kb.Instance.AddFile(c.Request.Context(), req.FileID, upsertOptions) - if err != nil { - // Update status to error and return error response + // Update status to error config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) - - errorResp := &response.ErrorResponse{ - Code: response.ErrServerError.Code, - ErrorDescription: "Failed to add file: " + err.Error(), - } - response.RespondWithError(c, response.StatusInternalServerError, errorResp) - return + return fmt.Errorf("failed to add file: %w", err) } // Update status to completed after successful processing @@ -86,6 +106,22 @@ func AddFile(c *gin.Context) { log.Error("Failed to update document status to completed: %v", err) } + 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", @@ -97,6 +133,39 @@ func AddFile(c *gin.Context) { 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 + } + + // Process the request + addFileWithRequest(c, &req) +} + // AddFileAsync adds file to a collection asynchronously func AddFileAsync(c *gin.Context) { var req AddFileRequest @@ -106,8 +175,23 @@ func AddFileAsync(c *gin.Context) { 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 := validateRequest(c, &req); err != nil { + if err := req.Validate(); err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } @@ -123,6 +207,12 @@ func AddFileAsync(c *gin.Context) { return } - // Handle async processing - handleAsync(c, AddFile) + // Handle async processing with parsed request + // Use context.Background() for async operations to avoid Gin context expiration + handleAsyncWithRequest(c, &req, func(ctx context.Context, r *AddFileRequest) { + err := AddFileProcess(ctx, r) + if err != nil { + log.Error("Async file processing failed: %v", err) + } + }) } diff --git a/openapi/kb/addtext.go b/openapi/kb/addtext.go index 5a4ee52c..a054fdd7 100644 --- a/openapi/kb/addtext.go +++ b/openapi/kb/addtext.go @@ -1,22 +1,98 @@ 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/kb" "github.com/yaoapp/yao/openapi/response" ) -// AddText adds text to a collection -func AddText(c *gin.Context) { +// ProcessAddTextRequest processes a text addition request with business logic only +// This function is Gin-agnostic and can be used for both sync and async operations +func ProcessAddTextRequest(ctx context.Context, req *AddTextRequest) error { // Check if kb.Instance is available - if !checkKBInstance(c) { - return + if kb.Instance == nil { + return fmt.Errorf("knowledge base not initialized") } + // Validate request + if err := req.Validate(); err != nil { + return err + } + + // Generate document ID if not provided + if req.DocID == "" { + req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + } + + // 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)), + } + + // 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) + + // 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() + 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 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) + } + + return nil +} + +// addTextWithRequest processes a text addition with pre-parsed request +func addTextWithRequest(c *gin.Context, req *AddTextRequest) { // Prepare request and database data - req, documentData, err := PrepareAddText(c) + _, documentData, err := PrepareAddText(c, req) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, @@ -87,12 +163,60 @@ func AddText(c *gin.Context) { response.RespondWithSuccess(c, response.StatusCreated, result) } +// AddText adds text to a collection +func AddText(c *gin.Context) { + var req AddTextRequest + + // 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 + } + + // Process the request + addTextWithRequest(c, &req) +} + // AddTextAsync adds text to a collection asynchronously func AddTextAsync(c *gin.Context) { var req AddTextRequest + // 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 := validateRequest(c, &req); err != nil { + if err := req.Validate(); err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } @@ -107,6 +231,11 @@ func AddTextAsync(c *gin.Context) { return } - // Handle async processing - handleAsync(c, AddText) + // Handle async processing with parsed request + handleAsyncWithRequest(c, &req, func(ctx context.Context, r *AddTextRequest) { + err := ProcessAddTextRequest(ctx, r) + if err != nil { + log.Error("Async text processing failed: %v", err) + } + }) } diff --git a/openapi/kb/addurl.go b/openapi/kb/addurl.go index d1c1c9b0..c5da3db5 100644 --- a/openapi/kb/addurl.go +++ b/openapi/kb/addurl.go @@ -1,6 +1,8 @@ package kb import ( + "context" + "github.com/gin-gonic/gin" "github.com/yaoapp/kun/log" "github.com/yaoapp/kun/maps" @@ -8,15 +10,10 @@ import ( "github.com/yaoapp/yao/openapi/response" ) -// AddURL adds a URL to a collection -func AddURL(c *gin.Context) { - // Check if kb.Instance is available - if !checkKBInstance(c) { - return - } - +// addURLWithRequest processes a URL addition with pre-parsed request +func addURLWithRequest(c *gin.Context, req *AddURLRequest) { // Prepare request and database data - req, documentData, err := PrepareAddURL(c) + _, documentData, err := PrepareAddURL(c, req) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, @@ -88,12 +85,60 @@ func AddURL(c *gin.Context) { response.RespondWithSuccess(c, response.StatusCreated, result) } +// AddURL adds a URL to a collection +func AddURL(c *gin.Context) { + var req AddURLRequest + + // 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 + } + + // Process the request + addURLWithRequest(c, &req) +} + // AddURLAsync adds a URL to a collection asynchronously func AddURLAsync(c *gin.Context) { var req AddURLRequest + // 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 := validateRequest(c, &req); err != nil { + if err := req.Validate(); err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } @@ -108,6 +153,9 @@ func AddURLAsync(c *gin.Context) { return } - // Handle async processing - handleAsync(c, AddURL) + // Handle async processing with parsed request + handleAsyncWithRequest(c, &req, func(ctx context.Context, r *AddURLRequest) { + // Temporary placeholder - would need ProcessAddURLRequest function + log.Info("Async URL processing placeholder for: %s", r.URL) + }) } diff --git a/openapi/kb/document.go b/openapi/kb/document.go index 2c37d250..9bb62183 100644 --- a/openapi/kb/document.go +++ b/openapi/kb/document.go @@ -1,6 +1,7 @@ package kb import ( + "context" "net/http" "github.com/gin-gonic/gin" @@ -67,31 +68,6 @@ type Validator interface { Validate() error } -// validateRequest validates a request by parsing JSON and calling Validate() -func validateRequest[T Validator](c *gin.Context, req T) error { - // 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 err - } - - // 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 err - } - - return nil -} - // checkKBInstance checks if kb.Instance is available func checkKBInstance(c *gin.Context) bool { if kb.Instance == nil { @@ -157,12 +133,13 @@ func validateFileAndGetPath(c *gin.Context, req *AddFileRequest) (string, string return path, contentType, nil } -// handleAsync handles async processing for any handler function -func handleAsync(c *gin.Context, syncHandler func(*gin.Context)) { +// handleAsyncWithRequest handles async processing for handlers that need parsed request data +func handleAsyncWithRequest[T any](c *gin.Context, req T, handler func(context.Context, T)) { jobid := uuid.New().String() // temporary solution to handle async operations ( TODO: use job queue ) - go func() { syncHandler(c) }() + // Use context.Background() to avoid Gin context expiration in async operations + go func() { handler(context.Background(), req) }() response.RespondWithSuccess(c, response.StatusCreated, gin.H{"job_id": jobid}) } diff --git a/openapi/kb/utils.go b/openapi/kb/utils.go index c0f83b8b..cea8817a 100644 --- a/openapi/kb/utils.go +++ b/openapi/kb/utils.go @@ -80,16 +80,14 @@ func PrepareCreateCollection(c *gin.Context) (*CreateCollectionRequest, map[stri } // PrepareAddFile prepares AddFile request and database data -func PrepareAddFile(c *gin.Context) (*AddFileRequest, map[string]interface{}, error) { - var req AddFileRequest - - // Parse and validate request - if err := validateRequest(c, &req); err != nil { +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) + path, contentType, err := validateFileAndGetPath(c, req) if err != nil { return nil, nil, err } @@ -120,15 +118,13 @@ func PrepareAddFile(c *gin.Context) (*AddFileRequest, map[string]interface{}, er req.BaseUpsertRequest.AddBaseFields(data) addContextFields(c, data) - return &req, data, nil + return req, data, nil } // PrepareAddText prepares AddText request and database data -func PrepareAddText(c *gin.Context) (*AddTextRequest, map[string]interface{}, error) { - var req AddTextRequest - - // Parse and validate request - if err := validateRequest(c, &req); err != nil { +func PrepareAddText(c *gin.Context, req *AddTextRequest) (*AddTextRequest, map[string]interface{}, error) { + // Validate request + if err := req.Validate(); err != nil { return nil, nil, err } @@ -158,15 +154,13 @@ func PrepareAddText(c *gin.Context) (*AddTextRequest, map[string]interface{}, er req.BaseUpsertRequest.AddBaseFields(data) addContextFields(c, data) - return &req, data, nil + return req, data, nil } // PrepareAddURL prepares AddURL request and database data -func PrepareAddURL(c *gin.Context) (*AddURLRequest, map[string]interface{}, error) { - var req AddURLRequest - - // Parse and validate request - if err := validateRequest(c, &req); err != nil { +func PrepareAddURL(c *gin.Context, req *AddURLRequest) (*AddURLRequest, map[string]interface{}, error) { + // Validate request + if err := req.Validate(); err != nil { return nil, nil, err } @@ -196,7 +190,7 @@ func PrepareAddURL(c *gin.Context) (*AddURLRequest, map[string]interface{}, erro req.BaseUpsertRequest.AddBaseFields(data) addContextFields(c, data) - return &req, data, nil + return req, data, nil } // addContextFields adds context-specific fields like permissions, user info