From 96eef9e27d891623eb80d6b940a040097ea49be9 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 29 Aug 2025 17:34:09 +0800 Subject: [PATCH] Implement document and collection count updates in API - Added DocumentCount and UpdateDocumentCount methods to manage document counts in collections, enhancing metadata accuracy. - Introduced RemoveDocumentsByCollectionID method for bulk document removal, improving collection management. - Updated AddFileProcess, AddTextProcess, and AddURLProcess functions to include document and segment count updates after file operations. - Enhanced RemoveCollection function to report the number of documents removed during collection deletion. - Implemented segment count updates in RemoveSegments and RemoveSegmentsByDocID functions, ensuring accurate tracking of document segments. --- kb/types/collection.go | 65 +++++++++++++++++++++++++ kb/types/document.go | 49 +++++++++++++++++++ openapi/kb/addfile.go | 19 ++++++++ openapi/kb/addtext.go | 19 ++++++++ openapi/kb/addurl.go | 19 ++++++++ openapi/kb/collection.go | 25 ++++++++-- openapi/kb/document.go | 102 ++++++++++++++++++++++++++++++++++++++- openapi/kb/segment.go | 30 ++++++++++++ 8 files changed, 322 insertions(+), 6 deletions(-) diff --git a/kb/types/collection.go b/kb/types/collection.go index 26e19d5f..8d59c351 100644 --- a/kb/types/collection.go +++ b/kb/types/collection.go @@ -5,6 +5,7 @@ import ( "github.com/yaoapp/gou/model" "github.com/yaoapp/kun/maps" + "github.com/yaoapp/xun/dbal" ) // SearchCollections searches collections with pagination @@ -108,3 +109,67 @@ func (c *Config) RemoveCollection(collectionID string) error { _, err := mod.DeleteWhere(param) return err } + +// DocumentCount returns the number of documents in a collection +func (c *Config) DocumentCount(collectionID string) (int, error) { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return 0, fmt.Errorf("document model not found: %s", modelName) + } + + // Use dbal.Raw to count documents in the collection + param := model.QueryParam{ + Select: []interface{}{dbal.Raw("COUNT(*) as count")}, + Wheres: []model.QueryWhere{ + {Column: "collection_id", Value: collectionID}, + }, + } + + result, err := mod.Get(param) + if err != nil { + return 0, fmt.Errorf("failed to count documents: %w", err) + } + + if len(result) == 0 { + return 0, nil + } + + // Extract count from result + countValue, exists := result[0]["count"] + if !exists { + return 0, fmt.Errorf("count field not found in result") + } + + // Convert to int + switch v := countValue.(type) { + case int: + return v, nil + case int64: + return int(v), nil + case float64: + return int(v), nil + default: + return 0, fmt.Errorf("unexpected count type: %T", v) + } +} + +// UpdateDocumentCount updates the document_count field in collection metadata +func (c *Config) UpdateDocumentCount(collectionID string) error { + // Get current document count + count, err := c.DocumentCount(collectionID) + if err != nil { + return fmt.Errorf("failed to get document count: %w", err) + } + + // Update collection metadata with the new count + data := maps.MapStrAny{ + "document_count": count, + } + + return c.UpdateCollection(collectionID, data) +} diff --git a/kb/types/document.go b/kb/types/document.go index cbe8faec..b161d092 100644 --- a/kb/types/document.go +++ b/kb/types/document.go @@ -108,3 +108,52 @@ func (c *Config) RemoveDocument(documentID string) error { _, err := mod.DeleteWhere(param) return err } + +// RemoveDocumentsByCollectionID removes all documents belonging to a collection +func (c *Config) RemoveDocumentsByCollectionID(collectionID string) error { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("document model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "collection_id", Value: collectionID}, + }, + } + + _, err := mod.DeleteWhere(param) + return err +} + +// UpdateSegmentCount updates the segment_count field for a document +func (c *Config) UpdateSegmentCount(documentID string, count int) error { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("document model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "document_id", Value: documentID}, + }, + Limit: 1, + } + + data := maps.MapStrAny{ + "segment_count": count, + } + + _, err := mod.UpdateWhere(param, data) + return err +} diff --git a/openapi/kb/addfile.go b/openapi/kb/addfile.go index a1ee6ff2..c9fc223c 100644 --- a/openapi/kb/addfile.go +++ b/openapi/kb/addfile.go @@ -112,6 +112,25 @@ func AddFileProcess(ctx context.Context, req *AddFileRequest, jobID ...string) e 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 + if err := config.UpdateDocumentCount(req.CollectionID); 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 } diff --git a/openapi/kb/addtext.go b/openapi/kb/addtext.go index e253e2aa..2a2d8fef 100644 --- a/openapi/kb/addtext.go +++ b/openapi/kb/addtext.go @@ -91,6 +91,25 @@ func AddTextProcess(ctx context.Context, req *AddTextRequest, jobID ...string) e 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 + if err := config.UpdateDocumentCount(req.CollectionID); 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 } diff --git a/openapi/kb/addurl.go b/openapi/kb/addurl.go index f9661cfb..5085fde7 100644 --- a/openapi/kb/addurl.go +++ b/openapi/kb/addurl.go @@ -90,6 +90,25 @@ func AddURLProcess(ctx context.Context, req *AddURLRequest, jobID ...string) err 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 + if err := config.UpdateDocumentCount(req.CollectionID); 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 } diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index 4ffe3510..ff3efa01 100644 --- a/openapi/kb/collection.go +++ b/openapi/kb/collection.go @@ -145,17 +145,34 @@ func RemoveCollection(c *gin.Context) { return } - // Remove collection from database after successful GraphRag removal + // Remove collection and all its documents from database after successful GraphRag removal + documentsRemoved := 0 if config, err := kb.GetConfig(); err == nil { + // First, count documents in this collection (for reporting) + if count, err := config.DocumentCount(collectionID); err == nil { + documentsRemoved = count + } + + // Remove all documents belonging to this collection + if err := config.RemoveDocumentsByCollectionID(collectionID); err != nil { + log.Error("Failed to remove documents from collection %s: %v", collectionID, err) + } else { + log.Info("Removed %d documents from collection %s", documentsRemoved, collectionID) + } + + // Then remove the collection itself if err := config.RemoveCollection(collectionID); err != nil { log.Error("Failed to remove collection from database: %v", err) + } else { + log.Info("Successfully removed collection %s and %d documents", collectionID, documentsRemoved) } } successData := gin.H{ - "message": "Collection removed successfully", - "collection_id": collectionID, - "removed": removed, + "message": "Collection removed successfully", + "collection_id": collectionID, + "removed": removed, + "documents_removed": documentsRemoved, } response.RespondWithSuccess(c, response.StatusOK, successData) } diff --git a/openapi/kb/document.go b/openapi/kb/document.go index 08f53b96..becd67d7 100644 --- a/openapi/kb/document.go +++ b/openapi/kb/document.go @@ -360,8 +360,106 @@ func GetDocument(c *gin.Context) { // RemoveDocs removes documents by IDs func RemoveDocs(c *gin.Context) { - // TODO: Implement remove documents logic - c.JSON(http.StatusOK, gin.H{"message": "Documents removed"}) + // Check if kb.Instance is available + if !checkKBInstance(c) { + return + } + + // Parse document_ids from query parameter (comma-separated string) + docIDsParam := strings.TrimSpace(c.Query("document_ids")) + if docIDsParam == "" { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "document_ids query parameter is required", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Split comma-separated document IDs + docIDs := strings.Split(docIDsParam, ",") + var validDocIDs []string + for _, id := range docIDs { + id = strings.TrimSpace(id) + if id != "" { + validDocIDs = append(validDocIDs, id) + } + } + + if len(validDocIDs) == 0 { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "No valid document IDs provided", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Get KB config for database operations + config, err := kb.GetConfig() + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to get KB config: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Remove documents using GraphRAG + deletedCount, err := kb.Instance.RemoveDocs(c.Request.Context(), validDocIDs) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to remove documents: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Also remove documents from the database and track collections to update + dbDeletedCount := 0 + collectionsToUpdate := make(map[string]bool) // Track unique collection IDs + + for _, docID := range validDocIDs { + // Get document info before deletion to track collection + if docInfo, err := config.FindDocument(docID, model.QueryParam{ + Select: []interface{}{"collection_id"}, + }); err == nil && docInfo != nil { + if collectionID, ok := docInfo["collection_id"].(string); ok && collectionID != "" { + collectionsToUpdate[collectionID] = true + } + } + + if err := config.RemoveDocument(docID); err != nil { + // Log the error but don't fail the entire operation + // since the document was already removed from GraphRAG + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to remove document from database: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + dbDeletedCount++ + } + + // Update document counts for affected collections + for collectionID := range collectionsToUpdate { + if err := config.UpdateDocumentCount(collectionID); 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) + } + } + + // Return success response with deletion count + c.JSON(http.StatusOK, gin.H{ + "message": "Documents removed successfully", + "deleted_count": deletedCount, + "requested_count": len(validDocIDs), + "db_deleted_count": dbDeletedCount, + }) } // Validator interface for request validation diff --git a/openapi/kb/segment.go b/openapi/kb/segment.go index aed1e31e..b99b3180 100644 --- a/openapi/kb/segment.go +++ b/openapi/kb/segment.go @@ -333,6 +333,22 @@ func RemoveSegments(c *gin.Context) { return } + // Update segment count for the document if segments were removed + if removedCount > 0 { + // Get KB config for database operations + config, err := kb.GetConfig() + if err == nil { + // Get current segment count and update document + if segmentCount, err := kb.Instance.SegmentCount(c.Request.Context(), docID); err == nil { + if err := config.UpdateSegmentCount(docID, segmentCount); err != nil { + // Log error but don't fail the operation + // TODO: Add proper logging + // log.Error("Failed to update segment count for document %s: %v", docID, err) + } + } + } + } + // Return success response result := gin.H{ "message": "Segments removed successfully", @@ -377,6 +393,20 @@ func RemoveSegmentsByDocID(c *gin.Context) { return } + // Update segment count for the document (should be 0 after removing all segments) + if removedCount > 0 { + // Get KB config for database operations + config, err := kb.GetConfig() + if err == nil { + // After removing all segments, count should be 0 + if err := config.UpdateSegmentCount(docID, 0); err != nil { + // Log error but don't fail the operation + // TODO: Add proper logging + // log.Error("Failed to update segment count for document %s: %v", docID, err) + } + } + } + // Return success response result := gin.H{ "message": "Segments removed successfully",