generalization of deleting rows for each dataset (marketplace manufacturer seller)

src/api/models/embedding_models.py

@@ -42,9 +42,10 @@ class DeleteEmbeddingRequest(BaseModel):
     dataset_name: str
 
 
-class DeleteRowsRequest(BaseModel):
+class DeleteByColumnRequest(BaseModel):
     dataset_name: str
-    product_types_to_delete: List[str]
+    key_column: str
+    keys_to_delete: List[str]
 
 
 # Request model for the /embed endpoint

src/api/services/huggingface_service.py

@@ -102,11 +102,11 @@ class HuggingFaceService:
             raise DatasetDeleteError(f"Failed to delete dataset: {e}")
 
 
-    async def delete_rows_from_dataset(self, dataset_name: str, product_types_to_delete: List[str]):
+    async def delete_rows_from_dataset(self, dataset_name: str, key_column: str, keys_to_delete: List[str]):
         """
-        Loads a dataset, filters out rows based on a list of product types, and pushes it back.
+        Loads a dataset, filters out rows based on a list of keys in a specific column, and pushes it back.
         """
-        if not product_types_to_delete:
+        if not keys_to_delete:
             return
 
         # Step 1: Load the existing dataset
@@ -114,10 +114,10 @@ class HuggingFaceService:
         dataset = await self.read_dataset(dataset_name)
 
         # Step 2 : Filter the dataset to EXCLUDE the rows with the given product_types
-        logger.info(f"Filtering out product_types: {product_types_to_delete}")
+        logger.info(f"Filtering out rows where column {key_column} is in {keys_to_delete}")
         initial_row_count = len(dataset)
 
-        filtered_dataset = dataset.filter(lambda product: product['product_type'] not in product_types_to_delete)
+        filtered_dataset = dataset.filter(lambda element: element[key_column] not in keys_to_delete)
 
         final_row_count = len(filtered_dataset)
         logger.info(f"{initial_row_count - final_row_count} rows deleted.")

src/main.py

@@ -7,9 +7,9 @@ from typing import List, Dict
 from datasets import Dataset
 from src.api.models.embedding_models import (
     CreateEmbeddingRequest,
+    DeleteByColumnRequest,
     ReadEmbeddingRequest,
     UpdateEmbeddingRequest,
-    DeleteRowsRequest,
     DeleteEmbeddingRequest,
     EmbedRequest,
     SearchEmbeddingRequest,
@@ -231,21 +231,22 @@ async def delete_embeddings(
         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
 
 
-@app.post("/delete_rows")
-async def delete_rows(
-    request: DeleteRowsRequest,
+@app.post("/delete_rows_by_key")
+async def delete_rows_by_key(
+    request: DeleteByColumnRequest,
     huggingface_service: HuggingFaceService = Depends(get_huggingface_service)
 ):
     """
-    Deletes specific rows from a Hugging Face dataset based on their product_types"""
+    Deletes specific rows from a Hugging Face dataset based on a key column and values."""
     try:
         await huggingface_service.delete_rows_from_dataset(
-            request.dataset_name, request.product_types_to_delete
+            request.dataset_name, request.key_column, request.keys_to_delete
         )
         return {
             "message": "Rows deleted succesfully from dataset.",
             "dataset_name": request.dataset_name,
-            "deleted_product_types": request.product_types_to_delete,
+            "key_column": request.key_column,
+            "deleted_keys": request.keys_to_delete,
         }
     except DatasetNotFoundError as e:
         raise HTTPException(status_code=404, detail=str(e))