Merge pull request #416 from ollama-webui/dev

refac

Dockerfile

@@ -10,7 +10,7 @@ RUN npm ci
 COPY . .
 RUN npm run build
 
-FROM python:3.11-bookworm as base
+FROM python:3.11-slim-bookworm as base
 
 ENV ENV=prod
 
@@ -28,7 +28,7 @@ WORKDIR /app/backend
 
 COPY ./backend/requirements.txt ./requirements.txt
 RUN pip3 install -r requirements.txt
-RUN python -c "from sentence_transformers import SentenceTransformer; model = SentenceTransformer('all-MiniLM-L6-v2')"
+# RUN python -c "from sentence_transformers import SentenceTransformer; model = SentenceTransformer('all-MiniLM-L6-v2')"
 
 COPY ./backend .
 

backend/apps/rag/main.py

@@ -11,9 +11,14 @@ from fastapi import (
 from fastapi.middleware.cors import CORSMiddleware
 import os, shutil
 
-from chromadb.utils import embedding_functions
+# from chromadb.utils import embedding_functions
 
-from langchain_community.document_loaders import WebBaseLoader, TextLoader, PyPDFLoader
+from langchain_community.document_loaders import (
+    WebBaseLoader,
+    TextLoader,
+    PyPDFLoader,
+    CSVLoader,
+)
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.chains import RetrievalQA
@@ -23,15 +28,16 @@ from pydantic import BaseModel
 from typing import Optional
 
 import uuid
+import time
 
-
+from utils.misc import calculate_sha256
 from utils.utils import get_current_user
 from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
 from constants import ERROR_MESSAGES
 
-EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
-    model_name=EMBED_MODEL
-)
+# EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
+#     model_name=EMBED_MODEL
+# )
 
 app = FastAPI()
 
@@ -64,9 +70,7 @@ def store_data_in_vector_db(data, collection_name) -> bool:
     metadatas = [doc.metadata for doc in docs]
 
     try:
-        collection = CHROMA_CLIENT.create_collection(
-            name=collection_name, embedding_function=EMBEDDING_FUNC
-        )
+        collection = CHROMA_CLIENT.create_collection(name=collection_name)
 
         collection.add(
             documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts]
@@ -125,14 +129,13 @@ def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
 
 @app.post("/doc")
 def store_doc(
-    collection_name: str = Form(...),
+    collection_name: Optional[str] = Form(None),
     file: UploadFile = File(...),
     user=Depends(get_current_user),
 ):
     # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
-    file.filename = f"{collection_name}-{file.filename}"
 
-    if file.content_type not in ["application/pdf", "text/plain"]:
+    if file.content_type not in ["application/pdf", "text/plain", "text/csv"]:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
@@ -146,10 +149,17 @@ def store_doc(
             f.write(contents)
             f.close()
 
+        f = open(file_path, "rb")
+        if collection_name == None:
+            collection_name = calculate_sha256(f)[:63]
+        f.close()
+
         if file.content_type == "application/pdf":
             loader = PyPDFLoader(file_path)
         elif file.content_type == "text/plain":
             loader = TextLoader(file_path)
+        elif file.content_type == "text/csv":
+            loader = CSVLoader(file_path)
 
         data = loader.load()
         result = store_data_in_vector_db(data, collection_name)
@@ -181,7 +191,7 @@ def reset_vector_db(user=Depends(get_current_user)):
 
 
 @app.get("/reset")
-def reset(user=Depends(get_current_user)):
+def reset(user=Depends(get_current_user)) -> bool:
     if user.role == "admin":
         folder = f"{UPLOAD_DIR}"
         for filename in os.listdir(folder):
@@ -199,7 +209,7 @@ def reset(user=Depends(get_current_user)):
         except Exception as e:
             print(e)
 
-        return {"status": True}
+        return True
     else:
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,

src/lib/apis/rag/index.ts

@@ -103,3 +103,29 @@ export const queryVectorDB = async (
 
 	return res;
 };
+
+export const resetVectorDB = async (token: string) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/reset`, {
+		method: 'GET',
+		headers: {
+			Accept: 'application/json',
+			authorization: `Bearer ${token}`
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};

src/lib/components/chat/MessageInput.svelte

@@ -91,6 +91,26 @@
 		}
 	};
 
+	const uploadDoc = async (file) => {
+		console.log(file);
+
+		const doc = {
+			type: 'doc',
+			name: file.name,
+			collection_name: '',
+			upload_status: false,
+			error: ''
+		};
+
+		files = [...files, doc];
+		const res = await uploadDocToVectorDB(localStorage.token, '', file);
+
+		if (res) {
+			doc.upload_status = true;
+			files = files;
+		}
+	};
+
 	onMount(() => {
 		const dropZone = document.querySelector('body');
 
@@ -122,21 +142,8 @@
 					const file = inputFiles[0];
 					if (['image/gif', 'image/jpeg', 'image/png'].includes(file['type'])) {
 						reader.readAsDataURL(file);
-					} else if (['application/pdf', 'text/plain'].includes(file['type'])) {
-						console.log(file);
-						const hash = (await calculateSHA256(file)).substring(0, 63);
-						const res = await uploadDocToVectorDB(localStorage.token, hash, file);
-
-						if (res) {
-							files = [
-								...files,
-								{
-									type: 'doc',
-									name: file.name,
-									collection_name: res.collection_name
-								}
-							];
-						}
+					} else if (['application/pdf', 'text/plain', 'text/csv'].includes(file['type'])) {
+						uploadDoc(file);
 					} else {
 						toast.error(`Unsupported File Type '${file['type']}'.`);
 					}
@@ -241,22 +248,9 @@
 							const file = inputFiles[0];
 							if (['image/gif', 'image/jpeg', 'image/png'].includes(file['type'])) {
 								reader.readAsDataURL(file);
-							} else if (['application/pdf', 'text/plain'].includes(file['type'])) {
-								console.log(file);
-								const hash = (await calculateSHA256(file)).substring(0, 63);
-								const res = await uploadDocToVectorDB(localStorage.token, hash, file);
-
-								if (res) {
-									files = [
-										...files,
-										{
-											type: 'doc',
-											name: file.name,
-											collection_name: res.collection_name
-										}
-									];
-									filesInputElement.value = '';
-								}
+							} else if (['application/pdf', 'text/plain', 'text/csv'].includes(file['type'])) {
+								uploadDoc(file);
+								filesInputElement.value = '';
 							} else {
 								toast.error(`Unsupported File Type '${file['type']}'.`);
 								inputFiles = null;
@@ -283,21 +277,65 @@
 											class="h-16 w-[15rem] flex items-center space-x-3 px-2.5 dark:bg-gray-600 rounded-xl border border-gray-200 dark:border-none"
 										>
 											<div class="p-2.5 bg-red-400 text-white rounded-lg">
-												<svg
-													xmlns="http://www.w3.org/2000/svg"
-													viewBox="0 0 24 24"
-													fill="currentColor"
-													class="w-6 h-6"
-												>
-													<path
-														fill-rule="evenodd"
-														d="M5.625 1.5c-1.036 0-1.875.84-1.875 1.875v17.25c0 1.035.84 1.875 1.875 1.875h12.75c1.035 0 1.875-.84 1.875-1.875V12.75A3.75 3.75 0 0 0 16.5 9h-1.875a1.875 1.875 0 0 1-1.875-1.875V5.25A3.75 3.75 0 0 0 9 1.5H5.625ZM7.5 15a.75.75 0 0 1 .75-.75h7.5a.75.75 0 0 1 0 1.5h-7.5A.75.75 0 0 1 7.5 15Zm.75 2.25a.75.75 0 0 0 0 1.5H12a.75.75 0 0 0 0-1.5H8.25Z"
-														clip-rule="evenodd"
-													/>
-													<path
-														d="M12.971 1.816A5.23 5.23 0 0 1 14.25 5.25v1.875c0 .207.168.375.375.375H16.5a5.23 5.23 0 0 1 3.434 1.279 9.768 9.768 0 0 0-6.963-6.963Z"
-													/>
-												</svg>
+												{#if file.upload_status}
+													<svg
+														xmlns="http://www.w3.org/2000/svg"
+														viewBox="0 0 24 24"
+														fill="currentColor"
+														class="w-6 h-6"
+													>
+														<path
+															fill-rule="evenodd"
+															d="M5.625 1.5c-1.036 0-1.875.84-1.875 1.875v17.25c0 1.035.84 1.875 1.875 1.875h12.75c1.035 0 1.875-.84 1.875-1.875V12.75A3.75 3.75 0 0 0 16.5 9h-1.875a1.875 1.875 0 0 1-1.875-1.875V5.25A3.75 3.75 0 0 0 9 1.5H5.625ZM7.5 15a.75.75 0 0 1 .75-.75h7.5a.75.75 0 0 1 0 1.5h-7.5A.75.75 0 0 1 7.5 15Zm.75 2.25a.75.75 0 0 0 0 1.5H12a.75.75 0 0 0 0-1.5H8.25Z"
+															clip-rule="evenodd"
+														/>
+														<path
+															d="M12.971 1.816A5.23 5.23 0 0 1 14.25 5.25v1.875c0 .207.168.375.375.375H16.5a5.23 5.23 0 0 1 3.434 1.279 9.768 9.768 0 0 0-6.963-6.963Z"
+														/>
+													</svg>
+												{:else}
+													<svg
+														class=" w-6 h-6 translate-y-[0.5px]"
+														fill="currentColor"
+														viewBox="0 0 24 24"
+														xmlns="http://www.w3.org/2000/svg"
+														><style>
+															.spinner_qM83 {
+																animation: spinner_8HQG 1.05s infinite;
+															}
+															.spinner_oXPr {
+																animation-delay: 0.1s;
+															}
+															.spinner_ZTLf {
+																animation-delay: 0.2s;
+															}
+															@keyframes spinner_8HQG {
+																0%,
+																57.14% {
+																	animation-timing-function: cubic-bezier(0.33, 0.66, 0.66, 1);
+																	transform: translate(0);
+																}
+																28.57% {
+																	animation-timing-function: cubic-bezier(0.33, 0, 0.66, 0.33);
+																	transform: translateY(-6px);
+																}
+																100% {
+																	transform: translate(0);
+																}
+															}
+														</style><circle class="spinner_qM83" cx="4" cy="12" r="2.5" /><circle
+															class="spinner_qM83 spinner_oXPr"
+															cx="12"
+															cy="12"
+															r="2.5"
+														/><circle
+															class="spinner_qM83 spinner_ZTLf"
+															cx="20"
+															cy="12"
+															r="2.5"
+														/></svg
+													>
+												{/if}
 											</div>
 
 											<div class="flex flex-col justify-center -space-y-0.5">

src/lib/components/chat/SettingsModal.svelte

@@ -31,6 +31,7 @@
 		updateOpenAIKey,
 		updateOpenAIUrl
 	} from '$lib/apis/openai';
+	import { resetVectorDB } from '$lib/apis/rag';
 
 	export let show = false;
 
@@ -1829,6 +1830,40 @@
 									<div class=" self-center text-sm font-medium">Delete All Chats</div>
 								</button>
 							{/if}
+
+							{#if $user?.role === 'admin'}
+								<hr class=" dark:border-gray-700" />
+
+								<button
+									class=" flex rounded-md py-2 px-3.5 w-full hover:bg-gray-200 dark:hover:bg-gray-800 transition"
+									on:click={() => {
+										const res = resetVectorDB(localStorage.token).catch((error) => {
+											toast.error(error);
+											return null;
+										});
+
+										if (res) {
+											toast.success('Success');
+										}
+									}}
+								>
+									<div class=" self-center mr-3">
+										<svg
+											xmlns="http://www.w3.org/2000/svg"
+											viewBox="0 0 16 16"
+											fill="currentColor"
+											class="w-4 h-4"
+										>
+											<path
+												fill-rule="evenodd"
+												d="M3.5 2A1.5 1.5 0 0 0 2 3.5v9A1.5 1.5 0 0 0 3.5 14h9a1.5 1.5 0 0 0 1.5-1.5v-7A1.5 1.5 0 0 0 12.5 4H9.621a1.5 1.5 0 0 1-1.06-.44L7.439 2.44A1.5 1.5 0 0 0 6.38 2H3.5Zm6.75 7.75a.75.75 0 0 0 0-1.5h-4.5a.75.75 0 0 0 0 1.5h4.5Z"
+												clip-rule="evenodd"
+											/>
+										</svg>
+									</div>
+									<div class=" self-center text-sm font-medium">Reset Vector Storage</div>
+								</button>
+							{/if}
 						</div>
 					</div>
 				{:else if selectedTab === 'auth'}

src/routes/(app)/+page.svelte

@@ -124,6 +124,14 @@
 		} else if (messages.length != 0 && messages.at(-1).done != true) {
 			// Response not done
 			console.log('wait');
+		} else if (
+			files.length > 0 &&
+			files.filter((file) => file.upload_status === false).length > 0
+		) {
+			// Upload not done
+			toast.error(
+				`Oops! Hold tight! Your files are still in the processing oven. We're cooking them up to perfection. Please be patient and we'll let you know once they're ready.`
+			);
 		} else {
 			// Reset chat message textarea height
 			document.getElementById('chat-textarea').style.height = '';