From a2e0fbc9430205c98ff24a75ddecb342c792aa08 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 24 Dec 2024 16:49:32 -0700 Subject: [PATCH] refac: collection query status --- backend/open_webui/utils/middleware.py | 66 ++++++++++++++++++- src/lib/components/chat/Chat.svelte | 33 ---------- .../chat/Messages/ResponseMessage.svelte | 58 +++++++++------- 3 files changed, 100 insertions(+), 57 deletions(-) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 19f732105..3700fa3b3 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -416,9 +416,12 @@ async def process_chat_payload(request, form_data, metadata, user, model): form_data = apply_params_to_form_data(form_data, model) log.debug(f"form_data: {form_data}") + event_emitter = get_event_emitter(metadata) + event_call = get_event_call(metadata) + extra_params = { - "__event_emitter__": get_event_emitter(metadata), - "__event_call__": get_event_call(metadata), + "__event_emitter__": event_emitter, + "__event_call__": event_call, "__user__": { "id": user.id, "email": user.email, @@ -432,9 +435,52 @@ async def process_chat_payload(request, form_data, metadata, user, model): # Initialize events to store additional event to be sent to the client # Initialize contexts and citation models = request.app.state.MODELS + events = [] sources = [] + user_message = get_last_user_message(form_data["messages"]) + model_knowledge = model.get("info", {}).get("meta", {}).get("knowledge", False) + + if model_knowledge: + await event_emitter( + { + "type": "status", + "data": { + "action": "knowledge_search", + "query": user_message, + "done": False, + }, + } + ) + + knowledge_files = [] + for item in model_knowledge: + print(item) + if item.get("collection_name"): + knowledge_files.append( + { + "id": item.get("collection_name"), + "name": item.get("name"), + "legacy": True, + } + ) + elif item.get("collection_names"): + knowledge_files.append( + { + "name": item.get("name"), + "type": "collection", + "collection_names": item.get("collection_names"), + "legacy": True, + } + ) + else: + knowledge_files.append(item) + + files = form_data.get("files", []) + files.extend(knowledge_files) + form_data["files"] = files + try: form_data, flags = await chat_completion_filter_functions_handler( request, form_data, model, extra_params @@ -445,6 +491,9 @@ async def process_chat_payload(request, form_data, metadata, user, model): tool_ids = form_data.pop("tool_ids", None) files = form_data.pop("files", None) + # Remove files duplicates + files = list({json.dumps(f, sort_keys=True): f for f in files}.values()) + metadata = { **metadata, "tool_ids": tool_ids, @@ -522,6 +571,19 @@ async def process_chat_payload(request, form_data, metadata, user, model): if len(sources) > 0: events.append({"sources": sources}) + if model_knowledge: + await event_emitter( + { + "type": "status", + "data": { + "action": "knowledge_search", + "query": user_message, + "done": True, + "hidden": True, + }, + } + ) + return form_data, events diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index d6120d0ee..88f4aaf3c 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -1440,39 +1440,6 @@ const userMessage = history.messages[responseMessage.parentId]; let files = JSON.parse(JSON.stringify(chatFiles)); - if (model?.info?.meta?.knowledge ?? false) { - // Only initialize and add status if knowledge exists - responseMessage.statusHistory = [ - { - action: 'knowledge_search', - description: $i18n.t(`Searching Knowledge for "{{searchQuery}}"`, { - searchQuery: userMessage.content - }), - done: false - } - ]; - files.push( - ...model.info.meta.knowledge.map((item) => { - if (item?.collection_name) { - return { - id: item.collection_name, - name: item.name, - legacy: true - }; - } else if (item?.collection_names) { - return { - name: item.name, - type: 'collection', - collection_names: item.collection_names, - legacy: true - }; - } else { - return item; - } - }) - ); - history.messages[responseMessageId] = responseMessage; - } files.push( ...(userMessage?.files ?? []).filter((item) => ['doc', 'file', 'collection'].includes(item.type) diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte index ae056531f..b8ca4fe79 100644 --- a/src/lib/components/chat/Messages/ResponseMessage.svelte +++ b/src/lib/components/chat/Messages/ResponseMessage.svelte @@ -519,37 +519,51 @@ {@const status = ( message?.statusHistory ?? [...(message?.status ? [message?.status] : [])] ).at(-1)} -
- {#if status?.done === false} -
- -
- {/if} + {#if !status?.hidden} +
+ {#if status?.done === false} +
+ +
+ {/if} - {#if status?.action === 'web_search' && status?.urls} - + {#if status?.action === 'web_search' && status?.urls} + +
+
+ {status?.description} +
+
+ + {:else if status?.action === 'knowledge_search'}
+ {$i18n.t(`Searching Knowledge for "{{searchQuery}}"`, { + searchQuery: status.query + })} +
+
+ {:else} +
+
{status?.description}
- - {:else} -
-
- {status?.description} -
-
- {/if} -
+ {/if} +
+ {/if} {/if} {#if edit === true}