Merge pull request #1106 from trheyi/main

Refactor file, text, and URL addition handlers for improved request p…

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
+// AddFileProcess processes a file addition request with business logic only
-func AddFile(c *gin.Context) {
+// 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) {
+	if kb.Instance == nil {
-		return
+		return fmt.Errorf("knowledge base not initialized")
 	}
 
-	// Prepare request and database data
+	// Validate request
-	req, documentData, err := PrepareAddFile(c)
+	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{
+		return fmt.Errorf("failed to get local path: %w", err)
-			Code:             response.ErrInvalidRequest.Code,
+	}
-			ErrorDescription: err.Error(),
+
-		}
+	fileInfo, err := m.Info(ctx, req.FileID)
-		response.RespondWithError(c, response.StatusBadRequest, errorResp)
+	if err != nil {
-		return
+		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{
+		return fmt.Errorf("failed to get KB config: %w", err)
-			Code:             response.ErrServerError.Code,
-			ErrorDescription: "Failed to get KB config: " + err.Error(),
-		}
-		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
-		return
 	}
 
+	// 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{
+		return fmt.Errorf("failed to save document metadata: %w", err)
-			Code:             response.ErrServerError.Code,
-			ErrorDescription: "Failed to save document metadata: " + err.Error(),
-		}
-		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
-		return
 	}
 
 	// 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 {
+		if rollbackErr := config.RemoveDocument(req.DocID); rollbackErr != nil {
-			log.Error("Failed to rollback document database record: %v", err)
+			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
+		// Update status to error
-		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
 		config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()})
-
+		return fmt.Errorf("failed to add file: %w", err)
-		errorResp := &response.ErrorResponse{
-			Code:             response.ErrServerError.Code,
-			ErrorDescription: "Failed to add file: " + err.Error(),
-		}
-		response.RespondWithError(c, response.StatusInternalServerError, errorResp)
-		return
 	}
 
 	// 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
+	// Handle async processing with parsed request
-	handleAsync(c, AddFile)
+	// 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)
+		}
+	})
 }

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
+// ProcessAddTextRequest processes a text addition request with business logic only
-func AddText(c *gin.Context) {
+// 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) {
+	if kb.Instance == nil {
-		return
+		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
+	// Handle async processing with parsed request
-	handleAsync(c, AddText)
+	handleAsyncWithRequest(c, &req, func(ctx context.Context, r *AddTextRequest) {
+		err := ProcessAddTextRequest(ctx, r)
+		if err != nil {
+			log.Error("Async text processing failed: %v", err)
+		}
+	})
 }

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
+// addURLWithRequest processes a URL addition with pre-parsed request
-func AddURL(c *gin.Context) {
+func addURLWithRequest(c *gin.Context, req *AddURLRequest) {
-	// Check if kb.Instance is available
-	if !checkKBInstance(c) {
-		return
-	}
-
 	// 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
+	// Handle async processing with parsed request
-	handleAsync(c, AddURL)
+	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)
+	})
 }

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
+// handleAsyncWithRequest handles async processing for handlers that need parsed request data
-func handleAsync(c *gin.Context, syncHandler func(*gin.Context)) {
+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})
 }

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) {
+func PrepareAddFile(c *gin.Context, req *AddFileRequest) (*AddFileRequest, map[string]interface{}, error) {
-	var req AddFileRequest
+	// Validate request
-
+	if err := req.Validate(); err != nil {
-	// Parse and validate request
-	if err := validateRequest(c, &req); 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) {
+func PrepareAddText(c *gin.Context, req *AddTextRequest) (*AddTextRequest, map[string]interface{}, error) {
-	var req AddTextRequest
+	// Validate request
-
+	if err := req.Validate(); err != nil {
-	// Parse and validate request
-	if err := validateRequest(c, &req); 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) {
+func PrepareAddURL(c *gin.Context, req *AddURLRequest) (*AddURLRequest, map[string]interface{}, error) {
-	var req AddURLRequest
+	// Validate request
-
+	if err := req.Validate(); err != nil {
-	// Parse and validate request
-	if err := validateRequest(c, &req); 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