diff --git a/openapi/kb/addfile.go b/openapi/kb/addfile.go index c9fc223c..44f56c1b 100644 --- a/openapi/kb/addfile.go +++ b/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) diff --git a/openapi/kb/addtext.go b/openapi/kb/addtext.go index 2a2d8fef..8f268e4b 100644 --- a/openapi/kb/addtext.go +++ b/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) diff --git a/openapi/kb/addurl.go b/openapi/kb/addurl.go index 5085fde7..f9a44f04 100644 --- a/openapi/kb/addurl.go +++ b/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) diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index 7ff0184e..1c1fb84a 100644 --- a/openapi/kb/collection.go +++ b/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) } diff --git a/openapi/kb/document.go b/openapi/kb/document.go index 3ebc3263..14358645 100644 --- a/openapi/kb/document.go +++ b/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) diff --git a/openapi/kb/utils.go b/openapi/kb/utils.go index dff06bfb..aa84ecf9 100644 --- a/openapi/kb/utils.go +++ b/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 +}