refactor: reorganize batch processing utilities and update import paths across services

app/api/v1/admin.py

@@ -12,7 +12,7 @@ from typing import Optional, List, Tuple
 from pydantic import BaseModel
 from app.core.auth import verify_api_key, verify_app_key, get_admin_api_key
 from app.core.config import config, get_config
-from app.core.batch_tasks import create_task, get_task, expire_task
+from app.core.batch import create_task, get_task, expire_task
 from app.core.storage import get_storage, LocalStorage, RedisStorage, SQLStorage
 from app.core.exceptions import AppException
 from app.services.token.manager import get_token_manager

app/core/batch.py

@@ -1,13 +1,94 @@
 """
-Batch task manager for admin batch operations (SSE progress).
-"""
+Batch utilities.
 
-from __future__ import annotations
+- run_batch: generic batch concurrency runner
+- BatchTask: SSE task manager for admin batch operations
+"""
 
 import asyncio
 import time
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar
+
+from app.core.logger import logger
+
+T = TypeVar("T")
+
+
+async def run_batch(
+    items: List[str],
+    worker: Callable[[str], Awaitable[T]],
+    *,
+    max_concurrent: int = 10,
+    batch_size: int = 50,
+    task: Optional["BatchTask"] = None,
+    on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None,
+    should_cancel: Optional[Callable[[], bool]] = None,
+) -> Dict[str, Dict[str, Any]]:
+    """
+    分批并发执行，单项失败不影响整体
+
+    Args:
+        items: 待处理项列表
+        worker: 异步处理函数
+        max_concurrent: 最大并发数
+        batch_size: 每批大小
+
+    Returns:
+        {item: {"ok": bool, "data": ..., "error": ...}}
+    """
+    try:
+        max_concurrent = int(max_concurrent)
+    except Exception:
+        max_concurrent = 10
+    try:
+        batch_size = int(batch_size)
+    except Exception:
+        batch_size = 50
+
+    max_concurrent = max(1, max_concurrent)
+    batch_size = max(1, batch_size)
+
+    sem = asyncio.Semaphore(max_concurrent)
+
+    async def _one(item: str) -> tuple[str, dict]:
+        if (should_cancel and should_cancel()) or (task and task.cancelled):
+            return item, {"ok": False, "error": "cancelled", "cancelled": True}
+        async with sem:
+            try:
+                data = await worker(item)
+                result = {"ok": True, "data": data}
+                if task:
+                    task.record(True)
+                if on_item:
+                    try:
+                        await on_item(item, result)
+                    except Exception:
+                        pass
+                return item, result
+            except Exception as e:
+                logger.warning(f"Batch item failed: {item[:16]}... - {e}")
+                result = {"ok": False, "error": str(e)}
+                if task:
+                    task.record(False, error=str(e))
+                if on_item:
+                    try:
+                        await on_item(item, result)
+                    except Exception:
+                        pass
+                return item, result
+
+    results: Dict[str, dict] = {}
+
+    # 分批执行，避免一次性创建所有 task
+    for i in range(0, len(items), batch_size):
+        if (should_cancel and should_cancel()) or (task and task.cancelled):
+            break
+        chunk = items[i : i + batch_size]
+        pairs = await asyncio.gather(*(_one(x) for x in chunk))
+        results.update(dict(pairs))
+
+    return results
 
 
 class BatchTask:
@@ -150,3 +231,13 @@ def delete_task(task_id: str) -> None:
 async def expire_task(task_id: str, delay: int = 300) -> None:
     await asyncio.sleep(delay)
     delete_task(task_id)
+
+
+__all__ = [
+    "run_batch",
+    "BatchTask",
+    "create_task",
+    "get_task",
+    "delete_task",
+    "expire_task",
+]

app/services/grok/batch_services/assets.py

@@ -11,7 +11,7 @@ from app.core.config import get_config
 from app.core.logger import logger
 from app.services.reverse.assets_list import AssetsListReverse
 from app.services.reverse.assets_delete import AssetsDeleteReverse
-from app.services.grok.utils.batch import run_in_batches
+from app.core.batch import run_batch
 
 
 class BaseAssetsService:
@@ -149,7 +149,7 @@ class ListService(BaseAssetsService):
                 return {"detail": detail, "count": 0}
 
         try:
-            return await run_in_batches(
+            return await run_batch(
                 tokens,
                 _fetch_detail,
                 max_concurrent=max_concurrent,
@@ -219,7 +219,7 @@ class DeleteService(BaseAssetsService):
                 return {"status": "error", "error": str(e)}
 
         try:
-            return await run_in_batches(
+            return await run_batch(
                 tokens,
                 _clear_one,
                 max_concurrent=max_concurrent,

app/services/grok/batch_services/nsfw.py

@@ -12,7 +12,7 @@ from app.core.exceptions import UpstreamException
 from app.services.reverse.accept_tos import AcceptTosReverse
 from app.services.reverse.nsfw_mgmt import NsfwMgmtReverse
 from app.services.reverse.set_birth import SetBirthReverse
-from app.services.grok.utils.batch import run_in_batches
+from app.core.batch import run_batch
 
 
 class NSFWService:
@@ -85,7 +85,7 @@ class NSFWService:
                 logger.error(f"NSFW enable failed: {e}")
                 return {"success": False, "http_status": 0, "error": str(e)[:100]}
 
-        return await run_in_batches(
+        return await run_batch(
             tokens,
             _enable,
             max_concurrent=max_concurrent,

app/services/grok/batch_services/usage.py

@@ -10,7 +10,7 @@ from curl_cffi.requests import AsyncSession
 from app.core.logger import logger
 from app.core.config import get_config
 from app.services.reverse.rate_limits import RateLimitsReverse
-from app.services.grok.utils.batch import run_in_batches
+from app.core.batch import run_batch
 
 _USAGE_SEMAPHORE = asyncio.Semaphore(25)
 _USAGE_SEM_VALUE = 25
@@ -71,7 +71,7 @@ class UsageService:
         async def _refresh_one(t: str):
             return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False)
 
-        return await run_in_batches(
+        return await run_batch(
             tokens,
             _refresh_one,
             max_concurrent=max_concurrent,

app/services/grok/utils/batch.py

@@ -1,86 +0,0 @@
-"""
-批量执行工具
-
-提供分批并发、单项失败隔离的通用批量处理能力。
-"""
-
-import asyncio
-from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar
-
-from app.core.logger import logger
-
-T = TypeVar("T")
-
-
-async def run_in_batches(
-    items: List[str],
-    worker: Callable[[str], Awaitable[T]],
-    *,
-    max_concurrent: int = 10,
-    batch_size: int = 50,
-    on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None,
-    should_cancel: Optional[Callable[[], bool]] = None,
-) -> Dict[str, Dict[str, Any]]:
-    """
-    分批并发执行，单项失败不影响整体
-
-    Args:
-        items: 待处理项列表
-        worker: 异步处理函数
-        max_concurrent: 最大并发数
-        batch_size: 每批大小
-
-    Returns:
-        {item: {"ok": bool, "data": ..., "error": ...}}
-    """
-    try:
-        max_concurrent = int(max_concurrent)
-    except Exception:
-        max_concurrent = 10
-    try:
-        batch_size = int(batch_size)
-    except Exception:
-        batch_size = 50
-
-    max_concurrent = max(1, max_concurrent)
-    batch_size = max(1, batch_size)
-
-    sem = asyncio.Semaphore(max_concurrent)
-
-    async def _one(item: str) -> tuple[str, dict]:
-        if should_cancel and should_cancel():
-            return item, {"ok": False, "error": "cancelled", "cancelled": True}
-        async with sem:
-            try:
-                data = await worker(item)
-                result = {"ok": True, "data": data}
-                if on_item:
-                    try:
-                        await on_item(item, result)
-                    except Exception:
-                        pass
-                return item, result
-            except Exception as e:
-                logger.warning(f"Batch item failed: {item[:16]}... - {e}")
-                result = {"ok": False, "error": str(e)}
-                if on_item:
-                    try:
-                        await on_item(item, result)
-                    except Exception:
-                        pass
-                return item, result
-
-    results: Dict[str, dict] = {}
-
-    # 分批执行，避免一次性创建所有 task
-    for i in range(0, len(items), batch_size):
-        if should_cancel and should_cancel():
-            break
-        chunk = items[i : i + batch_size]
-        pairs = await asyncio.gather(*(_one(x) for x in chunk))
-        results.update(dict(pairs))
-
-    return results
-
-
-__all__ = ["run_in_batches"]