diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index e95de90a5..1d1efc5df 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -372,7 +372,11 @@ from open_webui.utils.auth import ( from open_webui.utils.oauth import OAuthManager from open_webui.utils.security_headers import SecurityHeadersMiddleware -from open_webui.tasks import stop_task, list_tasks # Import from tasks.py +from open_webui.tasks import ( + list_task_ids_by_chat_id, + stop_task, + list_tasks, +) # Import from tasks.py from open_webui.utils.redis import get_sentinels_from_env @@ -1196,7 +1200,7 @@ async def chat_action( @app.post("/api/tasks/stop/{task_id}") async def stop_task_endpoint(task_id: str, user=Depends(get_verified_user)): try: - result = await stop_task(task_id) # Use the function from tasks.py + result = await stop_task(task_id) return result except ValueError as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) @@ -1204,7 +1208,19 @@ async def stop_task_endpoint(task_id: str, user=Depends(get_verified_user)): @app.get("/api/tasks") async def list_tasks_endpoint(user=Depends(get_verified_user)): - return {"tasks": list_tasks()} # Use the function from tasks.py + return {"tasks": list_tasks()} + + +@app.get("/api/tasks/chat/{chat_id}") +async def list_tasks_by_chat_id_endpoint(chat_id: str, user=Depends(get_verified_user)): + chat = Chats.get_chat_by_id(chat_id) + if chat is None or chat.user_id != user.id: + return {"task_ids": []} + + task_ids = list_task_ids_by_chat_id(chat_id) + + print(f"Task IDs for chat {chat_id}: {task_ids}") + return {"task_ids": task_ids} ################################## diff --git a/backend/open_webui/tasks.py b/backend/open_webui/tasks.py index 2740ecb5a..e575e6885 100644 --- a/backend/open_webui/tasks.py +++ b/backend/open_webui/tasks.py @@ -5,16 +5,23 @@ from uuid import uuid4 # A dictionary to keep track of active tasks tasks: Dict[str, asyncio.Task] = {} +chat_tasks = {} -def cleanup_task(task_id: str): +def cleanup_task(task_id: str, id=None): """ Remove a completed or canceled task from the global `tasks` dictionary. """ 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(coroutine): + +def create_task(coroutine, id=None): """ Create a new asyncio task and add it to the global task dictionary. """ @@ -22,9 +29,15 @@ def create_task(coroutine): task = asyncio.create_task(coroutine) # Create the task # Add a done callback for cleanup - task.add_done_callback(lambda t: cleanup_task(task_id)) - + task.add_done_callback(lambda t: cleanup_task(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] + return task_id, task @@ -42,6 +55,13 @@ def list_tasks(): return list(tasks.keys()) +def list_task_ids_by_chat_id(id): + """ + List all tasks associated with a specific ID. + """ + return chat_tasks.get(id, []) + + async def stop_task(task_id: str): """ Cancel a running task and remove it from the global task list. diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 3eb9be3c8..0f3dc67f5 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -2245,7 +2245,9 @@ async def process_chat_response( await response.background() # background_tasks.add_task(post_response_handler, response, events) - task_id, _ = create_task(post_response_handler(response, events)) + task_id, _ = create_task( + post_response_handler(response, events), id=metadata["chat_id"] + ) return {"status": True, "task_id": task_id} else: diff --git a/src/lib/apis/index.ts b/src/lib/apis/index.ts index 30364fabf..0929e0a69 100644 --- a/src/lib/apis/index.ts +++ b/src/lib/apis/index.ts @@ -260,6 +260,38 @@ export const stopTask = async (token: string, id: string) => { return res; }; +export const getTaskIdsByChatId = async (token: string, chat_id: string) => { + let error = null; + + const res = await fetch(`${WEBUI_BASE_URL}/api/tasks/chat/${chat_id}`, { + method: 'GET', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + ...(token && { authorization: `Bearer ${token}` }) + } + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .catch((err) => { + console.log(err); + if ('detail' in err) { + error = err.detail; + } else { + error = err; + } + return null; + }); + + if (error) { + throw error; + } + + return res; +}; + export const getToolServerData = async (token: string, url: string) => { let error = null; diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index 98226f990..8a8dd6fb7 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -74,7 +74,8 @@ generateQueries, chatAction, generateMoACompletion, - stopTask + stopTask, + getTaskIdsByChatId } from '$lib/apis'; import { getTools } from '$lib/apis/tools'; @@ -825,7 +826,14 @@ } } - taskIds = chat?.task_ids ?? null; + const taskRes = await getTaskIdsByChatId(localStorage.token, $chatId).catch((error) => { + return null; + }); + + if (taskRes) { + taskIds = taskRes.task_ids; + } + await tick(); return true; @@ -1721,7 +1729,6 @@ taskIds = null; const responseMessage = history.messages[history.currentId]; - // Set all response messages to done for (const messageId of history.messages[responseMessage.parentId].childrenIds) { history.messages[messageId].done = true; @@ -2014,6 +2021,7 @@
{/if} - {#if !history.currentId || history.messages[history.currentId]?.done == true} - {#if prompt === '' && files.length === 0} -
- - - -
- {:else} -
- - - -
- {/if} - {:else} + {#if taskIds && taskIds.length > 0}
+ {:else if prompt === '' && files.length === 0} +
+ + + +
+ {:else} +
+ + + +
{/if}