yao/kb/api/addurl.go

package api

import (
	"context"
	"fmt"

	"github.com/yaoapp/gou/graphrag/utils"
	"github.com/yaoapp/kun/maps"
	"github.com/yaoapp/yao/job"
)

// AddURL adds a URL to a collection (sync)
func (instance *KBInstance) AddURL(ctx context.Context, params *AddURLParams) (*AddDocumentResult, error) {
	// Validate required parameters
	if params.CollectionID == "" {
		return nil, fmt.Errorf("collection_id is required")
	}
	if params.URL == "" {
		return nil, fmt.Errorf("url 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")
	}

	// Generate document ID if not provided
	docID := params.DocID
	if docID == "" {
		docID = utils.GenDocIDWithCollectionID(params.CollectionID)
	}

	// Create document record
	documentData := map[string]interface{}{
		"document_id":   docID,
		"collection_id": params.CollectionID,
		"name":          "URL Document",
		"type":          "url",
		"status":        "pending",
		"url":           params.URL,
	}

	// Use title from metadata if available
	if params.Metadata != nil {
		if title, ok := params.Metadata["title"].(string); ok && title != "" {
			documentData["name"] = title
		}
	}

	// 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 URL content
	params.DocID = docID // Ensure docID is set
	err = instance.processURL(ctx, docID, params)
	if err != nil {
		return nil, err
	}

	return &AddDocumentResult{
		Message:      "URL added successfully",
		CollectionID: params.CollectionID,
		DocID:        docID,
		URL:          params.URL,
	}, nil
}

// AddURLAsync adds a URL to a collection (async)
func (instance *KBInstance) AddURLAsync(ctx context.Context, params *AddURLParams) (*AddDocumentAsyncResult, error) {
	// Validate required parameters
	if params.CollectionID == "" {
		return nil, fmt.Errorf("collection_id is required")
	}
	if params.URL == "" {
		return nil, fmt.Errorf("url 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")
	}

	// Generate document ID if not provided
	docID := params.DocID
	if docID == "" {
		docID = utils.GenDocIDWithCollectionID(params.CollectionID)
	}

	// Get job options with defaults
	jobName, jobDescription, jobIcon, jobCategory := getJobOptions(params.Job,
		"Knowledge Base Web Content Processing",
		"Fetching and indexing web 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":          "URL Document",
		"type":          "url",
		"status":        "pending",
		"url":           params.URL,
		"job_id":        j.JobID,
	}

	// Use title from metadata if available
	if params.Metadata != nil {
		if title, ok := params.Metadata["title"].(string); ok && title != "" {
			documentData["name"] = title
		}
	}

	// 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 := &AddURLParams{
		CollectionID: params.CollectionID,
		URL:          params.URL,
		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.addurl", func(execCtx *job.ExecutionContext) error {
		return instance.processURL(execCtx.Ctx, asyncDocID, asyncParams)
	}, map[string]interface{}{
		"doc_id":        asyncDocID,
		"collection_id": params.CollectionID,
		"url":           params.URL,
	})
	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
}

// processURL processes URL content and updates the knowledge base
func (instance *KBInstance) processURL(ctx context.Context, docID string, params *AddURLParams) error {
	// Convert to UpsertOptions
	upsertOptions, err := instance.toUpsertOptions(docID, params.CollectionID, params.Locale, "", "", 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 URL to GraphRag
	_, err = instance.GraphRag.AddURL(ctx, params.URL, upsertOptions)
	if err != nil {
		instance.Config.UpdateDocument(docID, maps.MapStrAny{"status": "error", "error_message": err.Error()})
		return fmt.Errorf("failed to add URL: %w", err)
	}

	// Update status and segment count
	instance.updateDocumentAfterProcessing(ctx, docID, params.CollectionID)

	return nil
}