Merge pull request #1126 from trheyi/main

Refactor document and collection update processes to include GraphRag…

openapi/kb/addfile.go

@@ -124,8 +124,8 @@ func AddFileProcess(ctx context.Context, req *AddFileRequest, jobID ...string) e
 		}
 	}
 
-	// Update document count for the collection
-	if err := config.UpdateDocumentCount(req.CollectionID); err != nil {
+	// 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)

openapi/kb/addtext.go

@@ -103,8 +103,8 @@ func AddTextProcess(ctx context.Context, req *AddTextRequest, jobID ...string) e
 		}
 	}
 
-	// Update document count for the collection
-	if err := config.UpdateDocumentCount(req.CollectionID); err != nil {
+	// 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)

openapi/kb/addurl.go

@@ -102,8 +102,8 @@ func AddURLProcess(ctx context.Context, req *AddURLRequest, jobID ...string) err
 		}
 	}
 
-	// Update document count for the collection
-	if err := config.UpdateDocumentCount(req.CollectionID); err != nil {
+	// 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)

openapi/kb/collection.go

@@ -124,8 +124,8 @@ func CreateCollection(c *gin.Context) {
 		return
 	}
 
-	// Update status to active after successful creation
-	updateErr := config.UpdateCollection(req.ID, maps.MapStrAny{"status": "active"})
+	// Update status to active after successful creation and sync to GraphRag
+	updateErr := UpdateCollectionWithSync(req.ID, maps.MapStrAny{"status": "active"}, config)
 	if updateErr != nil {
 		log.Error("Failed to update collection status to active: %v", updateErr)
 	}
@@ -558,6 +558,7 @@ func UpdateCollectionMetadata(c *gin.Context) {
 	}
 
 	// Update collection metadata in database after successful GraphRag update
+	// Note: Only update database here, don't sync to GraphRag again (already done above)
 	if config, err := kb.GetConfig(); err == nil {
 		// Prepare update data from metadata
 		updateData := maps.MapStrAny{}
@@ -572,6 +573,7 @@ func UpdateCollectionMetadata(c *gin.Context) {
 		}
 
 		if len(updateData) > 0 {
+			// Only update database, don't sync to GraphRag again to avoid duplicate updates
 			if err := config.UpdateCollection(collectionID, updateData); err != nil {
 				log.Error("Failed to update collection in database: %v", err)
 			}

openapi/kb/document.go

@@ -444,9 +444,9 @@ func RemoveDocs(c *gin.Context) {
 		dbDeletedCount++
 	}
 
-	// Update document counts for affected collections
+	// Update document counts for affected collections and sync to GraphRag
 	for collectionID := range collectionsToUpdate {
-		if err := config.UpdateDocumentCount(collectionID); err != nil {
+		if err := UpdateDocumentCountWithSync(collectionID, config); err != nil {
 			// Log error but don't fail the operation
 			// TODO: Add proper logging
 			// log.Error("Failed to update document count for collection %s: %v", collectionID, err)

openapi/kb/utils.go

@@ -1,11 +1,15 @@
 package kb
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/gou/graphrag/utils"
+	"github.com/yaoapp/kun/maps"
 	"github.com/yaoapp/yao/attachment"
+	"github.com/yaoapp/yao/kb"
+	kbtypes "github.com/yaoapp/yao/kb/types"
 )
 
 // PrepareCreateCollection prepares CreateCollection request and database data
@@ -212,3 +216,65 @@ func addContextFields(c *gin.Context, data map[string]interface{}) {
 	// Example: data["permissions"] = c.Get("permissions")
 	// Example: data["tenant_id"] = c.GetString("tenant_id")
 }
+
+// UpdateCollectionWithSync updates collection metadata in database and syncs to GraphRag
+func UpdateCollectionWithSync(collectionID string, data maps.MapStrAny, config *kbtypes.Config) error {
+	// Create a copy of data for GraphRag to avoid contamination from database operations
+	// This is necessary because Gou's UpdateWhere method modifies the input data parameter
+	originalData := make(maps.MapStrAny)
+	for k, v := range data {
+		originalData[k] = v
+	}
+
+	// Update collection in database
+	if err := config.UpdateCollection(collectionID, data); err != nil {
+		return fmt.Errorf("failed to update collection in database: %w", err)
+	}
+
+	// Sync to GraphRag metadata if kb.Instance is available
+	if kb.Instance != nil {
+		// Convert the original (unmodified) data to map[string]interface{}
+		metadata := make(map[string]interface{})
+		for k, v := range originalData {
+			metadata[k] = v
+		}
+
+		// Update GraphRag metadata
+		ctx := context.Background()
+		if err := kb.Instance.UpdateCollectionMetadata(ctx, collectionID, metadata); err != nil {
+			return fmt.Errorf("failed to sync collection metadata to GraphRag: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// UpdateDocumentCountWithSync updates document count in database and syncs to GraphRag metadata
+func UpdateDocumentCountWithSync(collectionID string, config *kbtypes.Config) error {
+	// Update document count in database
+	if err := config.UpdateDocumentCount(collectionID); err != nil {
+		return fmt.Errorf("failed to update document count in database: %w", err)
+	}
+
+	// Sync to GraphRag metadata if kb.Instance is available
+	if kb.Instance != nil {
+		// Get the updated document count
+		count, err := config.DocumentCount(collectionID)
+		if err != nil {
+			return fmt.Errorf("failed to get document count for sync: %w", err)
+		}
+
+		// Prepare metadata for GraphRag
+		metadata := map[string]interface{}{
+			"document_count": count,
+		}
+
+		// Update GraphRag metadata
+		ctx := context.Background()
+		if err := kb.Instance.UpdateCollectionMetadata(ctx, collectionID, metadata); err != nil {
+			return fmt.Errorf("failed to sync document count to GraphRag: %w", err)
+		}
+	}
+
+	return nil
+}