diff --git a/Dockerfile b/Dockerfile index aa6c3d552..88cfc6b93 100644 --- a/Dockerfile +++ b/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 . diff --git a/backend/apps/rag/main.py b/backend/apps/rag/main.py index a59aac6cb..4cde679e8 100644 --- a/backend/apps/rag/main.py +++ b/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, diff --git a/src/lib/apis/rag/index.ts b/src/lib/apis/rag/index.ts index bafd03607..8a2b8cb4e 100644 --- a/src/lib/apis/rag/index.ts +++ b/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; +}; diff --git a/src/lib/components/chat/MessageInput.svelte b/src/lib/components/chat/MessageInput.svelte index 36511e595..c8baab5d0 100644 --- a/src/lib/components/chat/MessageInput.svelte +++ b/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" >
- - - - + {#if file.upload_status} + + + + + {:else} + + {/if}
diff --git a/src/lib/components/chat/SettingsModal.svelte b/src/lib/components/chat/SettingsModal.svelte index 976b6115a..03a025327 100644 --- a/src/lib/components/chat/SettingsModal.svelte +++ b/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 @@
Delete All Chats
{/if} + + {#if $user?.role === 'admin'} +
+ + + {/if}
{:else if selectedTab === 'auth'} diff --git a/src/routes/(app)/+page.svelte b/src/routes/(app)/+page.svelte index d433615c1..653169417 100644 --- a/src/routes/(app)/+page.svelte +++ b/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 = '';