refac

backend/open_webui/tasks.py

@@ -17,9 +17,9 @@ tasks: Dict[str, asyncio.Task] = {}
 item_tasks = {}
 
 
-REDIS_TASKS_KEY = f"{REDIS_KEY_PREFIX}:tasks"
-REDIS_ITEM_TASKS_KEY = f"{REDIS_KEY_PREFIX}:tasks:item"
-REDIS_PUBSUB_CHANNEL = f"{REDIS_KEY_PREFIX}:tasks:commands"
+REDIS_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks'
+REDIS_ITEM_TASKS_KEY = f'{REDIS_KEY_PREFIX}:tasks:item'
+REDIS_PUBSUB_CHANNEL = f'{REDIS_KEY_PREFIX}:tasks:commands'
 
 
 async def redis_task_command_listener(app):
@@ -28,17 +28,17 @@ async def redis_task_command_listener(app):
     await pubsub.subscribe(REDIS_PUBSUB_CHANNEL)
 
     async for message in pubsub.listen():
-        if message["type"] != "message":
+        if message['type'] != 'message':
             continue
         try:
-            command = json.loads(message["data"])
-            if command.get("action") == "stop":
-                task_id = command.get("task_id")
+            command = json.loads(message['data'])
+            if command.get('action') == 'stop':
+                task_id = command.get('task_id')
                 local_task = tasks.get(task_id)
                 if local_task:
                     local_task.cancel()
         except Exception as e:
-            log.exception(f"Error handling distributed task command: {e}")
+            log.exception(f'Error handling distributed task command: {e}')
 
 
 ### ------------------------------
@@ -48,9 +48,9 @@ async def redis_task_command_listener(app):
 
 async def redis_save_task(redis: Redis, task_id: str, item_id: Optional[str]):
     pipe = redis.pipeline()
-    pipe.hset(REDIS_TASKS_KEY, task_id, item_id or "")
+    pipe.hset(REDIS_TASKS_KEY, task_id, item_id or '')
     if item_id:
-        pipe.sadd(f"{REDIS_ITEM_TASKS_KEY}:{item_id}", task_id)
+        pipe.sadd(f'{REDIS_ITEM_TASKS_KEY}:{item_id}', task_id)
     await pipe.execute()
 
 
@@ -58,11 +58,11 @@ async def redis_cleanup_task(redis: Redis, task_id: str, item_id: Optional[str])
     pipe = redis.pipeline()
     pipe.hdel(REDIS_TASKS_KEY, task_id)
     if item_id:
-        pipe.srem(f"{REDIS_ITEM_TASKS_KEY}:{item_id}", task_id)
+        pipe.srem(f'{REDIS_ITEM_TASKS_KEY}:{item_id}', task_id)
         await pipe.execute()
         # Remove the set key entirely if no tasks remain for this item
-        if await redis.scard(f"{REDIS_ITEM_TASKS_KEY}:{item_id}") == 0:
-            await redis.delete(f"{REDIS_ITEM_TASKS_KEY}:{item_id}")
+        if await redis.scard(f'{REDIS_ITEM_TASKS_KEY}:{item_id}') == 0:
+            await redis.delete(f'{REDIS_ITEM_TASKS_KEY}:{item_id}')
     else:
         await pipe.execute()
 
@@ -72,15 +72,15 @@ async def redis_list_tasks(redis: Redis) -> List[str]:
 
 
 async def redis_list_item_tasks(redis: Redis, item_id: str) -> List[str]:
-    return list(await redis.smembers(f"{REDIS_ITEM_TASKS_KEY}:{item_id}"))
+    return list(await redis.smembers(f'{REDIS_ITEM_TASKS_KEY}:{item_id}'))
 
 
 async def redis_send_command(redis: Redis, command: dict):
     command_json = json.dumps(command)
     # RedisCluster doesn't expose publish() directly, but the
     # PUBLISH command broadcasts across all cluster nodes server-side.
-    if hasattr(redis, "nodes_manager"):
-        await redis.execute_command("PUBLISH", REDIS_PUBSUB_CHANNEL, command_json)
+    if hasattr(redis, 'nodes_manager'):
+        await redis.execute_command('PUBLISH', REDIS_PUBSUB_CHANNEL, command_json)
     else:
         await redis.publish(REDIS_PUBSUB_CHANNEL, command_json)
 
@@ -109,9 +109,7 @@ async def create_task(redis, coroutine, id=None):
     task = asyncio.create_task(coroutine)  # Create the task
 
     # Add a done callback for cleanup
-    task.add_done_callback(
-        lambda t: asyncio.create_task(cleanup_task(redis, task_id, id))
-    )
+    task.add_done_callback(lambda t: asyncio.create_task(cleanup_task(redis, task_id, id)))
     tasks[task_id] = task
 
     # If an ID is provided, associate the task with that ID
@@ -155,30 +153,30 @@ async def stop_task(redis, task_id: str):
         await redis_send_command(
             redis,
             {
-                "action": "stop",
-                "task_id": task_id,
+                'action': 'stop',
+                'task_id': task_id,
             },
         )
         # Always clean Redis directly — hdel/srem are idempotent, safe even
         # if the done_callback on the owning process also fires cleanup.
         await redis_cleanup_task(redis, task_id, item_id or None)
-        return {"status": True, "message": f"Task {task_id} stopped."}
+        return {'status': True, 'message': f'Task {task_id} stopped.'}
 
     task = tasks.pop(task_id, None)
     if not task:
-        return {"status": False, "message": f"Task with ID {task_id} not found."}
+        return {'status': False, 'message': f'Task with ID {task_id} not found.'}
 
     task.cancel()  # Request task cancellation
     try:
         await task  # Wait for the task to handle the cancellation
     except asyncio.CancelledError:
         # Task successfully canceled
-        return {"status": True, "message": f"Task {task_id} successfully stopped."}
+        return {'status': True, 'message': f'Task {task_id} successfully stopped.'}
 
     if task.cancelled() or task.done():
-        return {"status": True, "message": f"Task {task_id} successfully cancelled."}
+        return {'status': True, 'message': f'Task {task_id} successfully cancelled.'}
 
-    return {"status": True, "message": f"Cancellation requested for {task_id}."}
+    return {'status': True, 'message': f'Cancellation requested for {task_id}.'}
 
 
 async def stop_item_tasks(redis: Redis, item_id: str):
@@ -187,14 +185,14 @@ async def stop_item_tasks(redis: Redis, item_id: str):
     """
     task_ids = await list_task_ids_by_item_id(redis, item_id)
     if not task_ids:
-        return {"status": True, "message": f"No tasks found for item {item_id}."}
+        return {'status': True, 'message': f'No tasks found for item {item_id}.'}
 
     for task_id in task_ids:
         result = await stop_task(redis, task_id)
-        if not result["status"]:
+        if not result['status']:
             return result  # Return the first failure
 
-    return {"status": True, "message": f"All tasks for item {item_id} stopped."}
+    return {'status': True, 'message': f'All tasks for item {item_id} stopped.'}
 
 
 async def has_active_tasks(redis, chat_id: str) -> bool: