yao/kb/api/addfile.go

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
}