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

@@ -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)
+}

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
+}

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
 }
 

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
 }
 

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
 }
 

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)
 }

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

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",