Browse Source

refac: collection query status

Timothy Jaeryang Baek 4 months ago
parent
commit
a2e0fbc943

+ 64 - 2
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
 
 

+ 0 - 33
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)

+ 36 - 22
src/lib/components/chat/Messages/ResponseMessage.svelte

@@ -519,37 +519,51 @@
 							{@const status = (
 								message?.statusHistory ?? [...(message?.status ? [message?.status] : [])]
 							).at(-1)}
-							<div class="status-description flex items-center gap-2 py-0.5">
-								{#if status?.done === false}
-									<div class="">
-										<Spinner className="size-4" />
-									</div>
-								{/if}
+							{#if !status?.hidden}
+								<div class="status-description flex items-center gap-2 py-0.5">
+									{#if status?.done === false}
+										<div class="">
+											<Spinner className="size-4" />
+										</div>
+									{/if}
 
-								{#if status?.action === 'web_search' && status?.urls}
-									<WebSearchResults {status}>
+									{#if status?.action === 'web_search' && status?.urls}
+										<WebSearchResults {status}>
+											<div class="flex flex-col justify-center -space-y-0.5">
+												<div
+													class="{status?.done === false
+														? 'shimmer'
+														: ''} text-base line-clamp-1 text-wrap"
+												>
+													{status?.description}
+												</div>
+											</div>
+										</WebSearchResults>
+									{:else if status?.action === 'knowledge_search'}
 										<div class="flex flex-col justify-center -space-y-0.5">
 											<div
 												class="{status?.done === false
 													? 'shimmer'
-													: ''} text-base line-clamp-1 text-wrap"
+													: ''} text-gray-500 dark:text-gray-500 text-base line-clamp-1 text-wrap"
 											>
-												{status?.description}
+												{$i18n.t(`Searching Knowledge for "{{searchQuery}}"`, {
+													searchQuery: status.query
+												})}
 											</div>
 										</div>
-									</WebSearchResults>
-								{:else}
-									<div class="flex flex-col justify-center -space-y-0.5">
-										<div
-											class="{status?.done === false
-												? 'shimmer'
-												: ''} text-gray-500 dark:text-gray-500 text-base line-clamp-1 text-wrap"
-										>
-											{status?.description}
+									{:else}
+										<div class="flex flex-col justify-center -space-y-0.5">
+											<div
+												class="{status?.done === false
+													? 'shimmer'
+													: ''} text-gray-500 dark:text-gray-500 text-base line-clamp-1 text-wrap"
+											>
+												{status?.description}
+											</div>
 										</div>
-									</div>
-								{/if}
-							</div>
+									{/if}
+								</div>
+							{/if}
 						{/if}
 
 						{#if edit === true}