mirror of
https://github.com/open-webui/open-webui
synced 2025-06-09 07:56:42 +00:00
168 lines
5.1 KiB
Python
168 lines
5.1 KiB
Python
# tasks.py
|
|
import asyncio
|
|
from typing import Dict
|
|
from uuid import uuid4
|
|
import json
|
|
from redis import Redis
|
|
from fastapi import Request
|
|
from typing import Dict, List, Optional
|
|
|
|
# A dictionary to keep track of active tasks
|
|
tasks: Dict[str, asyncio.Task] = {}
|
|
chat_tasks = {}
|
|
|
|
|
|
REDIS_TASKS_KEY = "open-webui:tasks"
|
|
REDIS_CHAT_TASKS_KEY = "open-webui:tasks:chat"
|
|
REDIS_PUBSUB_CHANNEL = "open-webui:tasks:commands"
|
|
|
|
|
|
def is_redis(request: Request) -> bool:
|
|
# Called everywhere a request is available to check Redis
|
|
return hasattr(request.app.state, "redis") and isinstance(
|
|
request.app.state.redis, Redis
|
|
)
|
|
|
|
|
|
async def redis_task_command_listener(app):
|
|
redis: Redis = app.state.redis
|
|
pubsub = redis.pubsub()
|
|
await pubsub.subscribe(REDIS_PUBSUB_CHANNEL)
|
|
print("Subscribed to Redis task command channel")
|
|
|
|
async for message in pubsub.listen():
|
|
if message["type"] != "message":
|
|
continue
|
|
try:
|
|
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:
|
|
print(f"Error handling distributed task command: {e}")
|
|
|
|
|
|
### ------------------------------
|
|
### REDIS-ENABLED HANDLERS
|
|
### ------------------------------
|
|
|
|
|
|
def redis_save_task(redis: Redis, task_id: str, chat_id: Optional[str]):
|
|
pipe = redis.pipeline()
|
|
pipe.hset(REDIS_TASKS_KEY, task_id, chat_id or "")
|
|
if chat_id:
|
|
pipe.sadd(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}", task_id)
|
|
pipe.execute()
|
|
|
|
|
|
def redis_cleanup_task(redis: Redis, task_id: str, chat_id: Optional[str]):
|
|
pipe = redis.pipeline()
|
|
pipe.hdel(REDIS_TASKS_KEY, task_id)
|
|
if chat_id:
|
|
pipe.srem(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}", task_id)
|
|
if pipe.scard(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}").execute()[-1] == 0:
|
|
pipe.delete(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}") # Remove if empty set
|
|
pipe.execute()
|
|
|
|
|
|
def redis_list_tasks(redis: Redis) -> List[str]:
|
|
return list(redis.hkeys(REDIS_TASKS_KEY))
|
|
|
|
|
|
def redis_list_chat_tasks(redis: Redis, chat_id: str) -> List[str]:
|
|
return list(redis.smembers(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}"))
|
|
|
|
|
|
def redis_send_command(redis: Redis, command: dict):
|
|
redis.publish(REDIS_PUBSUB_CHANNEL, json.dumps(command))
|
|
|
|
|
|
def cleanup_task(request, task_id: str, id=None):
|
|
"""
|
|
Remove a completed or canceled task from the global `tasks` dictionary.
|
|
"""
|
|
if is_redis(request):
|
|
redis_cleanup_task(request.app.state.redis, task_id, id)
|
|
|
|
tasks.pop(task_id, None) # Remove the task if it exists
|
|
|
|
# If an ID is provided, remove the task from the chat_tasks dictionary
|
|
if id and task_id in chat_tasks.get(id, []):
|
|
chat_tasks[id].remove(task_id)
|
|
if not chat_tasks[id]: # If no tasks left for this ID, remove the entry
|
|
chat_tasks.pop(id, None)
|
|
|
|
|
|
def create_task(request, coroutine, id=None):
|
|
"""
|
|
Create a new asyncio task and add it to the global task dictionary.
|
|
"""
|
|
task_id = str(uuid4()) # Generate a unique ID for the task
|
|
task = asyncio.create_task(coroutine) # Create the task
|
|
|
|
# Add a done callback for cleanup
|
|
task.add_done_callback(lambda t: cleanup_task(request, task_id, id))
|
|
tasks[task_id] = task
|
|
|
|
# If an ID is provided, associate the task with that ID
|
|
if chat_tasks.get(id):
|
|
chat_tasks[id].append(task_id)
|
|
else:
|
|
chat_tasks[id] = [task_id]
|
|
|
|
if is_redis(request):
|
|
redis_save_task(request.app.state.redis, task_id, id)
|
|
|
|
return task_id, task
|
|
|
|
|
|
def list_tasks(request):
|
|
"""
|
|
List all currently active task IDs.
|
|
"""
|
|
if is_redis(request):
|
|
return redis_list_tasks(request.app.state.redis)
|
|
return list(tasks.keys())
|
|
|
|
|
|
def list_task_ids_by_chat_id(request, id):
|
|
"""
|
|
List all tasks associated with a specific ID.
|
|
"""
|
|
if is_redis(request):
|
|
return redis_list_chat_tasks(request.app.state.redis, id)
|
|
return chat_tasks.get(id, [])
|
|
|
|
|
|
async def stop_task(request, task_id: str):
|
|
"""
|
|
Cancel a running task and remove it from the global task list.
|
|
"""
|
|
if is_redis(request):
|
|
# PUBSUB: All instances check if they have this task, and stop if so.
|
|
redis_send_command(
|
|
request.app.state.redis,
|
|
{
|
|
"action": "stop",
|
|
"task_id": task_id,
|
|
},
|
|
)
|
|
# Optionally check if task_id still in Redis a few moments later for feedback?
|
|
return {"status": True, "message": f"Stop signal sent for {task_id}"}
|
|
|
|
task = tasks.get(task_id)
|
|
if not task:
|
|
raise ValueError(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
|
|
tasks.pop(task_id, None) # Remove it from the dictionary
|
|
return {"status": True, "message": f"Task {task_id} successfully stopped."}
|
|
|
|
return {"status": False, "message": f"Failed to stop task {task_id}."}
|