update batch import

app/services.py

@@ -81,6 +81,7 @@ class StickerService:
     def import_stickers(
         sticker_dataset: str,
         upload: bool = False,
+        save_to_milvus: bool = False,
         progress_callback: callable = None,
     ) -> List[str]:
         """导入表情包数据集
@@ -98,7 +99,7 @@ class StickerService:
             cache_folder = os.path.join(TEMP_DIR, 'cache/')
             img_folder = os.path.join(TEMP_DIR, 'images/')
             data_json_path = os.path.join(cache_folder, 'data.json')
-
+            stickers = []
             logger.info(f"start import dataset")
             # 解压数据集
             with zipfile.ZipFile(sticker_dataset, 'r') as zip_ref:
@@ -142,11 +143,16 @@ class StickerService:
                             image_filename = f"image_{random.randint(100000, 999999)}.png"
                             file_path = f"images/{image_filename}"
                             generate_temp_image(img_folder, image, image_filename)
-
-                            db.store_sticker("", description, "", file_path, image_hash)
-                            results.append(f"成功导入: {image_filename}")
+                            stickers.append({
+                                "title": "",
+                                "description": description,
+                                "tags": "",
+                                "file_path": file_path,
+                                "image_hash": image_hash
+                            })
                             
                             if progress_callback:
+                                results.append(f"成功导入: {image_filename}")
                                 progress_callback(file, "Imported")
                             
                         except Exception as e:
@@ -154,9 +160,15 @@ class StickerService:
                             results.append(f"处理失败 {file}: {str(e)}")
                             if progress_callback:
                                 progress_callback(file, f"Failed: {str(e)}")
+            
+
+            if save_to_milvus and len(stickers) > 0:
+                logger.info(f"save to milvus, {len(stickers)} stickers")
+                db.batch_store_stickers(stickers)
 
             # 上传到 HuggingFace
-            if upload:
+            if upload and len(stickers) > 0:
+                logger.info(f"upload to huggingface, {len(stickers)} stickers")
                 upload_folder_to_huggingface(img_folder)
 
             return results

app/ui.py

@@ -144,6 +144,7 @@ class StickerUI:
             
             with gr.Row():
                 self.upload_checkbox = gr.Checkbox(label="Upload to HuggingFace", value=False)
+                self.save_to_milvus_checkbox = gr.Checkbox(label="Save to Milvus", value=False)
             
             with gr.Row():
                 self.import_button.render()
@@ -154,12 +155,13 @@ class StickerUI:
                 fn=self._import_stickers_with_progress,
                 inputs=[
                     self.dataset_input,
-                    self.upload_checkbox
+                    self.upload_checkbox,
+                    self.save_to_milvus_checkbox
                 ],
                 outputs=self.import_output
             )
     
-    def _import_stickers_with_progress(self, dataset_path, upload, progress=gr.Progress()):
+    def _import_stickers_with_progress(self, dataset_path, upload, save_to_milvus, progress=gr.Progress()):
         """Import stickers with progress tracking."""
         try:
             # Count total files first
@@ -187,8 +189,8 @@ class StickerUI:
             results = sticker_service.import_stickers(
                 dataset_path,
                 upload=upload,
+                save_to_milvus=save_to_milvus,
                 progress_callback=update_progress,
-                total_files=total_files
             )
             
             return "\n".join(results)

main.py

@@ -72,7 +72,7 @@ async def api_delete_stickers(request: dict):
 
 
 @app.post("/api/import_dataset")
-async def api_import_dataset(file: UploadFile = File(...), upload: bool = False):
+async def api_import_dataset(file: UploadFile = File(...), upload: bool = False, save_to_milvus: bool = False):
     """Import sticker dataset from ZIP file"""
     try:
         # Create a temporary file to store the uploaded ZIP
@@ -82,7 +82,7 @@ async def api_import_dataset(file: UploadFile = File(...), upload: bool = False)
             temp_file_path = temp_file.name
         
         # Import the dataset
-        results = sticker_service.import_stickers(temp_file_path, upload)
+        results = sticker_service.import_stickers(temp_file_path, upload, save_to_milvus)
         
         # Clean up the temporary file
         os.unlink(temp_file_path)