yao/openapi/kb/addfile.go

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,
	})
}