浏览代码

enh: enable selecting individual files from collection

Timothy J. Baek 6 月之前
父节点
当前提交
5a96fcbeaf

+ 44 - 1
backend/open_webui/apps/webui/models/files.py

@@ -50,6 +50,14 @@ class FileModel(BaseModel):
 ####################
 
 
+class FileMeta(BaseModel):
+    name: Optional[str] = None
+    content_type: Optional[str] = None
+    size: Optional[int] = None
+
+    model_config = ConfigDict(extra="allow")
+
+
 class FileModelResponse(BaseModel):
     id: str
     user_id: str
@@ -57,8 +65,15 @@ class FileModelResponse(BaseModel):
 
     filename: str
     data: Optional[dict] = None
-    meta: dict
+    meta: FileMeta
+
+    created_at: int  # timestamp in epoch
+    updated_at: int  # timestamp in epoch
+
 
+class FileMetadataResponse(BaseModel):
+    id: str
+    meta: dict
     created_at: int  # timestamp in epoch
     updated_at: int  # timestamp in epoch
 
@@ -104,6 +119,19 @@ class FilesTable:
             except Exception:
                 return None
 
+    def get_file_metadata_by_id(self, id: str) -> Optional[FileMetadataResponse]:
+        with get_db() as db:
+            try:
+                file = db.get(File, id)
+                return FileMetadataResponse(
+                    id=file.id,
+                    meta=file.meta,
+                    created_at=file.created_at,
+                    updated_at=file.updated_at,
+                )
+            except Exception:
+                return None
+
     def get_files(self) -> list[FileModel]:
         with get_db() as db:
             return [FileModel.model_validate(file) for file in db.query(File).all()]
@@ -118,6 +146,21 @@ class FilesTable:
                 .all()
             ]
 
+    def get_file_metadatas_by_ids(self, ids: list[str]) -> list[FileMetadataResponse]:
+        with get_db() as db:
+            return [
+                FileMetadataResponse(
+                    id=file.id,
+                    meta=file.meta,
+                    created_at=file.created_at,
+                    updated_at=file.updated_at,
+                )
+                for file in db.query(File)
+                .filter(File.id.in_(ids))
+                .order_by(File.updated_at.desc())
+                .all()
+            ]
+
     def get_files_by_user_id(self, user_id: str) -> list[FileModel]:
         with get_db() as db:
             return [

+ 6 - 0
backend/open_webui/apps/webui/models/knowledge.py

@@ -6,6 +6,10 @@ import uuid
 
 from open_webui.apps.webui.internal.db import Base, get_db
 from open_webui.env import SRC_LOG_LEVELS
+
+from open_webui.apps.webui.models.files import FileMetadataResponse
+
+
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy import BigInteger, Column, String, Text, JSON
 
@@ -64,6 +68,8 @@ class KnowledgeResponse(BaseModel):
     created_at: int  # timestamp in epoch
     updated_at: int  # timestamp in epoch
 
+    files: Optional[list[FileMetadataResponse | dict]] = None
+
 
 class KnowledgeForm(BaseModel):
     name: str

+ 6 - 3
backend/open_webui/apps/webui/routers/files.py

@@ -213,7 +213,7 @@ async def update_file_data_content_by_id(
 ############################
 
 
-@router.get("/{id}/content", response_model=Optional[FileModel])
+@router.get("/{id}/content")
 async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
     file = Files.get_file_by_id(id)
 
@@ -239,7 +239,7 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
         )
 
 
-@router.get("/{id}/content/{file_name}", response_model=Optional[FileModel])
+@router.get("/{id}/content/{file_name}")
 async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
     file = Files.get_file_by_id(id)
 
@@ -251,7 +251,10 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
             # Check if the file already exists in the cache
             if file_path.is_file():
                 print(f"file_path: {file_path}")
-                return FileResponse(file_path)
+                headers = {
+                    "Content-Disposition": f'attachment; filename="{file.meta.get("name", file.filename)}"'
+                }
+                return FileResponse(file_path, headers=headers)
             else:
                 raise HTTPException(
                     status_code=status.HTTP_404_NOT_FOUND,

+ 6 - 1
backend/open_webui/apps/webui/routers/knowledge.py

@@ -48,7 +48,12 @@ async def get_knowledge_items(
             )
     else:
         return [
-            KnowledgeResponse(**knowledge.model_dump())
+            KnowledgeResponse(
+                **knowledge.model_dump(),
+                files=Files.get_file_metadatas_by_ids(
+                    knowledge.data.get("file_ids", []) if knowledge.data else []
+                ),
+            )
             for knowledge in Knowledges.get_knowledge_items()
         ]
 

+ 0 - 1
src/lib/components/chat/MessageInput/Commands.svelte

@@ -55,7 +55,6 @@
 				files = [
 					...files,
 					{
-						type: e?.detail?.meta?.document ? 'file' : 'collection',
 						...e.detail,
 						status: 'processed'
 					}

+ 125 - 33
src/lib/components/chat/MessageInput/Commands/Knowledge.svelte

@@ -2,6 +2,10 @@
 	import { toast } from 'svelte-sonner';
 	import Fuse from 'fuse.js';
 
+	import dayjs from 'dayjs';
+	import relativeTime from 'dayjs/plugin/relativeTime';
+	dayjs.extend(relativeTime);
+
 	import { createEventDispatcher, tick, getContext, onMount } from 'svelte';
 	import { removeLastWordFromString, isValidHttpUrl } from '$lib/utils';
 	import { knowledge } from '$lib/stores';
@@ -72,7 +76,13 @@
 	};
 
 	onMount(() => {
-		let legacy_documents = $knowledge.filter((item) => item?.meta?.document);
+		let legacy_documents = $knowledge
+			.filter((item) => item?.meta?.document)
+			.map((item) => ({
+				...item,
+				type: 'file'
+			}));
+
 		let legacy_collections =
 			legacy_documents.length > 0
 				? [
@@ -101,12 +111,44 @@
 					]
 				: [];
 
-		items = [...$knowledge, ...legacy_collections].map((item) => {
-			return {
+		let collections = $knowledge
+			.filter((item) => !item?.meta?.document)
+			.map((item) => ({
 				...item,
-				...(item?.legacy || item?.meta?.legacy || item?.meta?.document ? { legacy: true } : {})
-			};
-		});
+				type: 'collection'
+			}));
+		let collection_files =
+			$knowledge.length > 0
+				? [
+						...$knowledge
+							.reduce((a, item) => {
+								return [
+									...new Set([
+										...a,
+										...(item?.files ?? []).map((file) => ({
+											...file,
+											collection: { name: item.name, description: item.description }
+										}))
+									])
+								];
+							}, [])
+							.map((file) => ({
+								...file,
+								name: file?.meta?.name,
+								description: `${file?.collection?.name} - ${file?.collection?.description}`,
+								type: 'file'
+							}))
+					]
+				: [];
+
+		items = [...collections, ...collection_files, ...legacy_collections, ...legacy_documents].map(
+			(item) => {
+				return {
+					...item,
+					...(item?.legacy || item?.meta?.legacy || item?.meta?.document ? { legacy: true } : {})
+				};
+			}
+		);
 
 		fuse = new Fuse(items, {
 			keys: ['name', 'description']
@@ -126,7 +168,8 @@
 				<div class="m-1 overflow-y-auto p-1 rounded-r-xl space-y-0.5 scrollbar-hidden">
 					{#each filteredItems as item, idx}
 						<button
-							class=" px-3 py-1.5 rounded-xl w-full text-left {idx === selectedIdx
+							class=" px-3 py-1.5 rounded-xl w-full text-left flex justify-between items-center {idx ===
+							selectedIdx
 								? ' bg-gray-50 dark:bg-gray-850 dark:text-gray-100 selected-command-option-button'
 								: ''}"
 							type="button"
@@ -137,38 +180,87 @@
 							on:mousemove={() => {
 								selectedIdx = idx;
 							}}
-							on:focus={() => {}}
 						>
-							<div class=" font-medium text-black dark:text-gray-100 flex items-center gap-1">
-								{#if item.legacy}
-									<div
-										class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1"
-									>
-										Legacy
-									</div>
-								{:else if item?.meta?.document}
-									<div
-										class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1"
-									>
-										Document
-									</div>
-								{:else}
-									<div
-										class="bg-green-500/20 text-green-700 dark:text-green-200 rounded uppercase text-xs font-bold px-1"
-									>
-										Collection
-									</div>
-								{/if}
+							<div>
+								<div class=" font-medium text-black dark:text-gray-100 flex items-center gap-1">
+									{#if item.legacy}
+										<div
+											class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1 flex-shrink-0"
+										>
+											Legacy
+										</div>
+									{:else if item?.meta?.document}
+										<div
+											class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1 flex-shrink-0"
+										>
+											Document
+										</div>
+									{:else if item?.type === 'file'}
+										<div
+											class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1 flex-shrink-0"
+										>
+											File
+										</div>
+									{:else}
+										<div
+											class="bg-green-500/20 text-green-700 dark:text-green-200 rounded uppercase text-xs font-bold px-1 flex-shrink-0"
+										>
+											Collection
+										</div>
+									{/if}
 
-								<div class="line-clamp-1">
-									{item.name}
+									<div class="line-clamp-1">
+										{item?.name}
+									</div>
 								</div>
-							</div>
 
-							<div class=" text-xs text-gray-600 dark:text-gray-100 line-clamp-1">
-								{item?.description}
+								<div class=" text-xs text-gray-600 dark:text-gray-100 line-clamp-1">
+									{item?.description}
+								</div>
 							</div>
 						</button>
+
+						<!-- <div slot="content" class=" pl-2 pt-1 flex flex-col gap-0.5">
+								{#if !item.legacy && (item?.files ?? []).length > 0}
+									{#each item?.files ?? [] as file, fileIdx}
+										<button
+											class=" px-3 py-1.5 rounded-xl w-full text-left flex justify-between items-center hover:bg-gray-50 dark:hover:bg-gray-850 dark:hover:text-gray-100 selected-command-option-button"
+											type="button"
+											on:click={() => {
+												console.log(file);
+											}}
+											on:mousemove={() => {
+												selectedIdx = idx;
+											}}
+										>
+											<div>
+												<div
+													class=" font-medium text-black dark:text-gray-100 flex items-center gap-1"
+												>
+													<div
+														class="bg-gray-500/20 text-gray-700 dark:text-gray-200 rounded uppercase text-xs font-bold px-1 flex-shrink-0"
+													>
+														File
+													</div>
+
+													<div class="line-clamp-1">
+														{file?.meta?.name}
+													</div>
+												</div>
+
+												<div class=" text-xs text-gray-600 dark:text-gray-100 line-clamp-1">
+													{$i18n.t('Updated')}
+													{dayjs(file.updated_at * 1000).fromNow()}
+												</div>
+											</div>
+										</button>
+									{/each}
+								{:else}
+									<div class=" text-gray-500 text-xs mt-1 mb-2">
+										{$i18n.t('No files found.')}
+									</div>
+								{/if}
+							</div> -->
 					{/each}
 
 					{#if prompt