Merge pull request #333 from ollama-webui/rag

feat: RAG support

Dockerfile

@@ -10,7 +10,7 @@ RUN npm ci
 COPY . .
 RUN npm run build
 
-FROM python:3.11-slim-buster as base
+FROM python:3.11-bookworm as base
 
 ENV ENV=prod
 

README.md

@@ -33,6 +33,8 @@ Also check our sibling project, [OllamaHub](https://ollamahub.com/), where you c
 
 - ✒️🔢 **Full Markdown and LaTeX Support**: Elevate your LLM experience with comprehensive Markdown and LaTeX capabilities for enriched interaction.
 
+- 📚 **RAG Integration (Alpha)**: Immerse yourself in cutting-edge Retrieval Augmented Generation support, revolutionizing your chat experience by seamlessly incorporating document interactions. In its alpha phase, expect occasional issues as we actively refine and enhance this feature to ensure optimal performance and reliability.
+
 - 📜 **Prompt Preset Support**: Instantly access preset prompts using the '/' command in the chat input. Load predefined conversation starters effortlessly and expedite your interactions. Effortlessly import prompts through [OllamaHub](https://ollamahub.com/) integration.
 
 - 👍👎 **RLHF Annotation**: Empower your messages by rating them with thumbs up and thumbs down, facilitating the creation of datasets for Reinforcement Learning from Human Feedback (RLHF). Utilize your messages to train or fine-tune models, all while ensuring the confidentiality of locally saved data.
@@ -243,7 +245,6 @@ See [TROUBLESHOOTING.md](/TROUBLESHOOTING.md) for information on how to troubles
 
 Here are some exciting tasks on our roadmap:
 
-- 📚 **RAG Integration**: Experience first-class retrieval augmented generation support, enabling chat with your documents.
 - 🌐 **Web Browsing Capability**: Experience the convenience of seamlessly integrating web content directly into your chat. Easily browse and share information without leaving the conversation.
 - 🔄 **Function Calling**: Empower your interactions by running code directly within the chat. Execute functions and commands effortlessly, enhancing the functionality of your conversations.
 - ⚙️ **Custom Python Backend Actions**: Empower your Ollama Web UI by creating or downloading custom Python backend actions. Unleash the full potential of your web interface with tailored actions that suit your specific needs, enhancing functionality and versatility.

backend/.gitignore

@@ -5,4 +5,5 @@ uploads
 .ipynb_checkpoints
 *.db
 _test
-Pipfile
+Pipfile
+data/*

backend/apps/rag/main.py

@@ -0,0 +1,207 @@
+from fastapi import (
+    FastAPI,
+    Request,
+    Depends,
+    HTTPException,
+    status,
+    UploadFile,
+    File,
+    Form,
+)
+from fastapi.middleware.cors import CORSMiddleware
+import os, shutil
+
+from chromadb.utils import embedding_functions
+
+from langchain_community.document_loaders import WebBaseLoader, TextLoader, PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.chains import RetrievalQA
+
+
+from pydantic import BaseModel
+from typing import Optional
+
+import uuid
+
+
+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
+)
+
+app = FastAPI()
+
+origins = ["*"]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class CollectionNameForm(BaseModel):
+    collection_name: Optional[str] = "test"
+
+
+class StoreWebForm(CollectionNameForm):
+    url: str
+
+
+def store_data_in_vector_db(data, collection_name) -> bool:
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
+    )
+    docs = text_splitter.split_documents(data)
+
+    texts = [doc.page_content for doc in docs]
+    metadatas = [doc.metadata for doc in docs]
+
+    try:
+        collection = CHROMA_CLIENT.create_collection(
+            name=collection_name, embedding_function=EMBEDDING_FUNC
+        )
+
+        collection.add(
+            documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts]
+        )
+        return True
+    except Exception as e:
+        print(e)
+        if e.__class__.__name__ == "UniqueConstraintError":
+            return True
+
+        return False
+
+
+@app.get("/")
+async def get_status():
+    return {"status": True}
+
+
+@app.get("/query/{collection_name}")
+def query_collection(
+    collection_name: str,
+    query: str,
+    k: Optional[int] = 4,
+    user=Depends(get_current_user),
+):
+    try:
+        collection = CHROMA_CLIENT.get_collection(
+            name=collection_name,
+        )
+        result = collection.query(query_texts=[query], n_results=k)
+
+        return result
+    except Exception as e:
+        print(e)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.DEFAULT(e),
+        )
+
+
+@app.post("/web")
+def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
+    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
+    try:
+        loader = WebBaseLoader(form_data.url)
+        data = loader.load()
+        store_data_in_vector_db(data, form_data.collection_name)
+        return {"status": True, "collection_name": form_data.collection_name}
+    except Exception as e:
+        print(e)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.DEFAULT(e),
+        )
+
+
+@app.post("/doc")
+def store_doc(
+    collection_name: str = Form(...),
+    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"]:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
+        )
+
+    try:
+        filename = file.filename
+        file_path = f"{UPLOAD_DIR}/{filename}"
+        contents = file.file.read()
+        with open(file_path, "wb") as f:
+            f.write(contents)
+            f.close()
+
+        if file.content_type == "application/pdf":
+            loader = PyPDFLoader(file_path)
+        elif file.content_type == "text/plain":
+            loader = TextLoader(file_path)
+
+        data = loader.load()
+        result = store_data_in_vector_db(data, collection_name)
+
+        if result:
+            return {"status": True, "collection_name": collection_name}
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=ERROR_MESSAGES.DEFAULT(),
+            )
+    except Exception as e:
+        print(e)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.DEFAULT(e),
+        )
+
+
+@app.get("/reset/db")
+def reset_vector_db(user=Depends(get_current_user)):
+    if user.role == "admin":
+        CHROMA_CLIENT.reset()
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+        )
+
+
+@app.get("/reset")
+def reset(user=Depends(get_current_user)):
+    if user.role == "admin":
+        folder = f"{UPLOAD_DIR}"
+        for filename in os.listdir(folder):
+            file_path = os.path.join(folder, filename)
+            try:
+                if os.path.isfile(file_path) or os.path.islink(file_path):
+                    os.unlink(file_path)
+                elif os.path.isdir(file_path):
+                    shutil.rmtree(file_path)
+            except Exception as e:
+                print("Failed to delete %s. Reason: %s" % (file_path, e))
+
+        try:
+            CHROMA_CLIENT.reset()
+        except Exception as e:
+            print(e)
+
+        return {"status": True}
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
+        )

backend/config.py

@@ -1,14 +1,31 @@
 from dotenv import load_dotenv, find_dotenv
+import os
+
+
+import chromadb
+from chromadb import Settings
 
-from constants import ERROR_MESSAGES
 
 from secrets import token_bytes
 from base64 import b64encode
 
-import os
+from constants import ERROR_MESSAGES
+
+
+from pathlib import Path
 
 load_dotenv(find_dotenv("../.env"))
 
+
+####################################
+# File Upload
+####################################
+
+
+UPLOAD_DIR = "./data/uploads"
+Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
+
+
 ####################################
 # ENV (dev,test,prod)
 ####################################
@@ -19,8 +36,9 @@ ENV = os.environ.get("ENV", "dev")
 # OLLAMA_API_BASE_URL
 ####################################
 
-OLLAMA_API_BASE_URL = os.environ.get("OLLAMA_API_BASE_URL",
+OLLAMA_API_BASE_URL = os.environ.get(
-                                     "http://localhost:11434/api")
+    "OLLAMA_API_BASE_URL", "http://localhost:11434/api"
+)
 
 if ENV == "prod":
     if OLLAMA_API_BASE_URL == "/ollama/api":
@@ -56,3 +74,15 @@ WEBUI_JWT_SECRET_KEY = os.environ.get("WEBUI_JWT_SECRET_KEY", "t0p-s3cr3t")
 
 if WEBUI_AUTH and WEBUI_JWT_SECRET_KEY == "":
     raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND)
+
+####################################
+# RAG
+####################################
+
+CHROMA_DATA_PATH = "./data/vector_db"
+EMBED_MODEL = "all-MiniLM-L6-v2"
+CHROMA_CLIENT = chromadb.PersistentClient(
+    path=CHROMA_DATA_PATH, settings=Settings(allow_reset=True)
+)
+CHUNK_SIZE = 1500
+CHUNK_OVERLAP = 100

backend/constants.py

@@ -6,7 +6,6 @@ class MESSAGES(str, Enum):
 
 
 class ERROR_MESSAGES(str, Enum):
-
     def __str__(self) -> str:
         return super().__str__()
 
@@ -30,7 +29,12 @@ class ERROR_MESSAGES(str, Enum):
     UNAUTHORIZED = "401 Unauthorized"
     ACCESS_PROHIBITED = "You do not have permission to access this resource. Please contact your administrator for assistance."
     ACTION_PROHIBITED = (
-        "The requested action has been restricted as a security measure.")
+        "The requested action has been restricted as a security measure."
+    )
+
+    FILE_NOT_SENT = "FILE_NOT_SENT"
+    FILE_NOT_SUPPORTED = "Oops! It seems like the file format you're trying to upload is not supported. Please upload a file with a supported format (e.g., JPG, PNG, PDF, TXT) and try again."
+
     NOT_FOUND = "We could not find what you're looking for :/"
     USER_NOT_FOUND = "We could not find what you're looking for :/"
     API_KEY_NOT_FOUND = "Oops! It looks like there's a hiccup. The API key is missing. Please make sure to provide a valid API key to access this feature."

backend/main.py

@@ -1,3 +1,5 @@
+import time
+
 from fastapi import FastAPI, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi import HTTPException
@@ -5,16 +7,17 @@ from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
+
 from apps.ollama.main import app as ollama_app
 from apps.openai.main import app as openai_app
 
 from apps.web.main import app as webui_app
+from apps.rag.main import app as rag_app
 
-import time
+from config import ENV
 
 
 class SPAStaticFiles(StaticFiles):
-
     async def get_response(self, path: str, scope):
         try:
             return await super().get_response(path, scope)
@@ -25,7 +28,7 @@ class SPAStaticFiles(StaticFiles):
                 raise ex
 
 
-app = FastAPI()
+app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
 
 origins = ["*"]
 
@@ -49,9 +52,10 @@ async def check_url(request: Request, call_next):
 
 
 app.mount("/api/v1", webui_app)
+
 app.mount("/ollama/api", ollama_app)
 app.mount("/openai/api", openai_app)
+app.mount("/rag/api/v1", rag_app)
 
-app.mount("/",
+
-          SPAStaticFiles(directory="../build", html=True),
+app.mount("/", SPAStaticFiles(directory="../build", html=True), name="spa-static-files")
-          name="spa-static-files")

backend/requirements.txt

@@ -16,6 +16,13 @@ aiohttp
 peewee
 bcrypt
 
+
+langchain
+langchain-community
+chromadb
+sentence_transformers
+pypdf
+
 PyJWT
 pyjwt[crypto]
 

src/lib/apis/rag/index.ts

@@ -0,0 +1,105 @@
+import { RAG_API_BASE_URL } from '$lib/constants';
+
+export const uploadDocToVectorDB = async (token: string, collection_name: string, file: File) => {
+	const data = new FormData();
+	data.append('file', file);
+	data.append('collection_name', collection_name);
+
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/doc`, {
+		method: 'POST',
+		headers: {
+			Accept: 'application/json',
+			authorization: `Bearer ${token}`
+		},
+		body: data
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			error = err.detail;
+			console.log(err);
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
+export const uploadWebToVectorDB = async (token: string, collection_name: string, url: string) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/web`, {
+		method: 'POST',
+		headers: {
+			Accept: 'application/json',
+			'Content-Type': 'application/json',
+			authorization: `Bearer ${token}`
+		},
+		body: JSON.stringify({
+			url: url,
+			collection_name: collection_name
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			error = err.detail;
+			console.log(err);
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
+export const queryVectorDB = async (
+	token: string,
+	collection_name: string,
+	query: string,
+	k: number
+) => {
+	let error = null;
+	const searchParams = new URLSearchParams();
+
+	searchParams.set('query', query);
+	if (k) {
+		searchParams.set('k', k.toString());
+	}
+
+	const res = await fetch(
+		`${RAG_API_BASE_URL}/query/${collection_name}/?${searchParams.toString()}`,
+		{
+			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

@@ -2,10 +2,11 @@
 	import toast from 'svelte-french-toast';
 	import { onMount, tick } from 'svelte';
 	import { settings } from '$lib/stores';
-	import { findWordIndices } from '$lib/utils';
+	import { calculateSHA256, findWordIndices } from '$lib/utils';
 
 	import Prompts from './MessageInput/PromptCommands.svelte';
 	import Suggestions from './MessageInput/Suggestions.svelte';
+	import { uploadDocToVectorDB } from '$lib/apis/rag';
 
 	export let submitPrompt: Function;
 	export let stopResponse: Function;
@@ -98,7 +99,7 @@
 			dragged = true;
 		});
 
-		dropZone.addEventListener('drop', (e) => {
+		dropZone.addEventListener('drop', async (e) => {
 			e.preventDefault();
 			console.log(e);
 
@@ -115,14 +116,32 @@
 					];
 				};
 
-				if (
+				const inputFiles = e.dataTransfer?.files;
-					e.dataTransfer?.files &&
+
-					e.dataTransfer?.files.length > 0 &&
+				if (inputFiles && inputFiles.length > 0) {
-					['image/gif', 'image/jpeg', 'image/png'].includes(e.dataTransfer?.files[0]['type'])
+					const file = inputFiles[0];
-				) {
+					if (['image/gif', 'image/jpeg', 'image/png'].includes(file['type'])) {
-					reader.readAsDataURL(e.dataTransfer?.files[0]);
+						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 {
+						toast.error(`Unsupported File Type '${file['type']}'.`);
+					}
 				} else {
-					toast.error(`Unsupported File Type '${e.dataTransfer?.files[0]['type']}'.`);
+					toast.error(`File not found.`);
 				}
 			}
 
@@ -145,11 +164,11 @@
 		<div class="absolute rounded-xl w-full h-full backdrop-blur bg-gray-800/40 flex justify-center">
 			<div class="m-auto pt-64 flex flex-col justify-center">
 				<div class="max-w-md">
-					<div class="  text-center text-6xl mb-3">🏞️</div>
+					<div class="  text-center text-6xl mb-3">🗂️</div>
-					<div class="text-center dark:text-white text-2xl font-semibold z-50">Add Images</div>
+					<div class="text-center dark:text-white text-2xl font-semibold z-50">Add Files</div>
 
 					<div class=" mt-2 text-center text-sm dark:text-gray-200 w-full">
-						Drop any images here to add to the conversation
+						Drop any files/images here to add to the conversation
 					</div>
 				</div>
 			</div>
@@ -204,7 +223,7 @@
 					bind:files={inputFiles}
 					type="file"
 					hidden
-					on:change={() => {
+					on:change={async () => {
 						let reader = new FileReader();
 						reader.onload = (event) => {
 							files = [
@@ -218,15 +237,32 @@
 							filesInputElement.value = '';
 						};
 
-						if (
+						if (inputFiles && inputFiles.length > 0) {
-							inputFiles &&
+							const file = inputFiles[0];
-							inputFiles.length > 0 &&
+							if (['image/gif', 'image/jpeg', 'image/png'].includes(file['type'])) {
-							['image/gif', 'image/jpeg', 'image/png'].includes(inputFiles[0]['type'])
+								reader.readAsDataURL(file);
-						) {
+							} else if (['application/pdf', 'text/plain'].includes(file['type'])) {
-							reader.readAsDataURL(inputFiles[0]);
+								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 {
+								toast.error(`Unsupported File Type '${file['type']}'.`);
+								inputFiles = null;
+							}
 						} else {
-							toast.error(`Unsupported File Type '${inputFiles[0]['type']}'.`);
+							toast.error(`File not found.`);
-							inputFiles = null;
 						}
 					}}
 				/>
@@ -237,10 +273,42 @@
 					}}
 				>
 					{#if files.length > 0}
-						<div class="ml-2 mt-2 mb-1 flex space-x-2">
+						<div class="mx-2 mt-2 mb-1 flex flex-wrap gap-2">
 							{#each files as file, fileIdx}
 								<div class=" relative group">
-									<img src={file.url} alt="input" class=" h-16 w-16 rounded-xl object-cover" />
+									{#if file.type === 'image'}
+										<img src={file.url} alt="input" class=" h-16 w-16 rounded-xl object-cover" />
+									{:else if file.type === 'doc'}
+										<div
+											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>
+											</div>
+
+											<div class="flex flex-col justify-center -space-y-0.5">
+												<div class=" dark:text-gray-100 text-sm font-medium line-clamp-1">
+													{file.name}
+												</div>
+
+												<div class=" text-gray-500 text-sm">Document</div>
+											</div>
+										</div>
+									{/if}
 
 									<div class=" absolute -top-1 -right-1">
 										<button

src/lib/components/chat/Messages/UserMessage.svelte

@@ -53,11 +53,41 @@
 			class="prose chat-{message.role} w-full max-w-full dark:prose-invert prose-headings:my-0 prose-p:my-0 prose-p:-mb-4 prose-pre:my-0 prose-table:my-0 prose-blockquote:my-0 prose-img:my-0 prose-ul:-my-4 prose-ol:-my-4 prose-li:-my-3 prose-ul:-mb-6 prose-ol:-mb-6 prose-li:-mb-4 whitespace-pre-line"
 		>
 			{#if message.files}
-				<div class="my-3 w-full flex overflow-x-auto space-x-2">
+				<div class="my-2.5 w-full flex overflow-x-auto gap-2 flex-wrap">
 					{#each message.files as file}
 						<div>
 							{#if file.type === 'image'}
 								<img src={file.url} alt="input" class=" max-h-96 rounded-lg" draggable="false" />
+							{:else if file.type === 'doc'}
+								<div
+									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>
+									</div>
+
+									<div class="flex flex-col justify-center -space-y-0.5">
+										<div class=" dark:text-gray-100 text-sm font-medium line-clamp-1">
+											{file.name}
+										</div>
+
+										<div class=" text-gray-500 text-sm">Document</div>
+									</div>
+								</div>
 							{/if}
 						</div>
 					{/each}

src/lib/constants.ts

@@ -5,6 +5,7 @@ export const WEBUI_BASE_URL = dev ? `http://${location.hostname}:8080` : ``;
 export const WEBUI_API_BASE_URL = `${WEBUI_BASE_URL}/api/v1`;
 export const OLLAMA_API_BASE_URL = `${WEBUI_BASE_URL}/ollama/api`;
 export const OPENAI_API_BASE_URL = `${WEBUI_BASE_URL}/openai/api`;
+export const RAG_API_BASE_URL = `${WEBUI_BASE_URL}/rag/api/v1`;
 
 export const WEB_UI_VERSION = 'v1.0.0-alpha-static';
 

src/lib/utils/index.ts

@@ -127,3 +127,37 @@ export const findWordIndices = (text) => {
 
 	return matches;
 };
+
+export const calculateSHA256 = async (file) => {
+	// Create a FileReader to read the file asynchronously
+	const reader = new FileReader();
+
+	// Define a promise to handle the file reading
+	const readFile = new Promise((resolve, reject) => {
+		reader.onload = () => resolve(reader.result);
+		reader.onerror = reject;
+	});
+
+	// Read the file as an ArrayBuffer
+	reader.readAsArrayBuffer(file);
+
+	try {
+		// Wait for the FileReader to finish reading the file
+		const buffer = await readFile;
+
+		// Convert the ArrayBuffer to a Uint8Array
+		const uint8Array = new Uint8Array(buffer);
+
+		// Calculate the SHA-256 hash using Web Crypto API
+		const hashBuffer = await crypto.subtle.digest('SHA-256', uint8Array);
+
+		// Convert the hash to a hexadecimal string
+		const hashArray = Array.from(new Uint8Array(hashBuffer));
+		const hashHex = hashArray.map((byte) => byte.toString(16).padStart(2, '0')).join('');
+
+		return `${hashHex}`;
+	} catch (error) {
+		console.error('Error calculating SHA-256 hash:', error);
+		throw error;
+	}
+};

src/lib/utils/rag/index.ts

@@ -0,0 +1,20 @@
+export const RAGTemplate = (context: string, query: string) => {
+	let template = `Use the following context as your learned knowledge, inside <context></context> XML tags.
+	<context>
+	  [context]
+	</context>
+	
+	When answer to user:
+	- If you don't know, just say that you don't know.
+	- If you don't know when you are not sure, ask for clarification.
+	Avoid mentioning that you obtained the information from the context.
+	And answer according to the language of the user's question.
+			
+	Given the context information, answer the query.
+	Query: [query]`;
+
+	template = template.replace(/\[context\]/g, context);
+	template = template.replace(/\[query\]/g, query);
+
+	return template;
+};

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

@@ -7,16 +7,18 @@
 	import { page } from '$app/stores';
 
 	import { models, modelfiles, user, settings, chats, chatId, config } from '$lib/stores';
+	import { copyToClipboard, splitStream } from '$lib/utils';
 
 	import { generateChatCompletion, generateTitle } from '$lib/apis/ollama';
-	import { copyToClipboard, splitStream } from '$lib/utils';
+	import { createNewChat, getChatList, updateChatById } from '$lib/apis/chats';
+	import { queryVectorDB } from '$lib/apis/rag';
+	import { generateOpenAIChatCompletion } from '$lib/apis/openai';
 
 	import MessageInput from '$lib/components/chat/MessageInput.svelte';
 	import Messages from '$lib/components/chat/Messages.svelte';
 	import ModelSelector from '$lib/components/chat/ModelSelector.svelte';
 	import Navbar from '$lib/components/layout/Navbar.svelte';
-	import { createNewChat, getChatList, updateChatById } from '$lib/apis/chats';
+	import { RAGTemplate } from '$lib/utils/rag';
-	import { generateOpenAIChatCompletion } from '$lib/apis/openai';
 
 	let stopResponseFlag = false;
 	let autoScroll = true;
@@ -113,8 +115,108 @@
 	// Ollama functions
 	//////////////////////////
 
+	const submitPrompt = async (userPrompt) => {
+		console.log('submitPrompt', $chatId);
+
+		if (selectedModels.includes('')) {
+			toast.error('Model not selected');
+		} else if (messages.length != 0 && messages.at(-1).done != true) {
+			// Response not done
+			console.log('wait');
+		} else {
+			// Reset chat message textarea height
+			document.getElementById('chat-textarea').style.height = '';
+
+			// Create user message
+			let userMessageId = uuidv4();
+			let userMessage = {
+				id: userMessageId,
+				parentId: messages.length !== 0 ? messages.at(-1).id : null,
+				childrenIds: [],
+				role: 'user',
+				content: userPrompt,
+				files: files.length > 0 ? files : undefined
+			};
+
+			// Add message to history and Set currentId to messageId
+			history.messages[userMessageId] = userMessage;
+			history.currentId = userMessageId;
+
+			// Append messageId to childrenIds of parent message
+			if (messages.length !== 0) {
+				history.messages[messages.at(-1).id].childrenIds.push(userMessageId);
+			}
+
+			// Wait until history/message have been updated
+			await tick();
+
+			// Create new chat if only one message in messages
+			if (messages.length == 1) {
+				if ($settings.saveChatHistory ?? true) {
+					chat = await createNewChat(localStorage.token, {
+						id: $chatId,
+						title: 'New Chat',
+						models: selectedModels,
+						system: $settings.system ?? undefined,
+						options: {
+							...($settings.options ?? {})
+						},
+						messages: messages,
+						history: history,
+						timestamp: Date.now()
+					});
+					await chats.set(await getChatList(localStorage.token));
+					await chatId.set(chat.id);
+				} else {
+					await chatId.set('local');
+				}
+				await tick();
+			}
+
+			// Reset chat input textarea
+			prompt = '';
+			files = [];
+
+			// Send prompt
+			await sendPrompt(userPrompt, userMessageId);
+		}
+	};
+
 	const sendPrompt = async (prompt, parentId) => {
 		const _chatId = JSON.parse(JSON.stringify($chatId));
+
+		const docs = messages
+			.filter((message) => message?.files ?? null)
+			.map((message) => message.files.filter((item) => item.type === 'doc'))
+			.flat(1);
+
+		console.log(docs);
+		if (docs.length > 0) {
+			const query = history.messages[parentId].content;
+
+			let relevantContexts = await Promise.all(
+				docs.map(async (doc) => {
+					return await queryVectorDB(localStorage.token, doc.collection_name, query, 4).catch(
+						(error) => {
+							console.log(error);
+							return null;
+						}
+					);
+				})
+			);
+			relevantContexts = relevantContexts.filter((context) => context);
+
+			const contextString = relevantContexts.reduce((a, context, i, arr) => {
+				return `${a}${context.documents.join(' ')}\n`;
+			}, '');
+
+			console.log(contextString);
+
+			history.messages[parentId].raContent = RAGTemplate(contextString, query);
+			history.messages[parentId].contexts = relevantContexts;
+			await tick();
+		}
+
 		await Promise.all(
 			selectedModels.map(async (model) => {
 				console.log(model);
@@ -177,7 +279,7 @@
 				.filter((message) => message)
 				.map((message) => ({
 					role: message.role,
-					content: message.content,
+					content: message?.raContent ?? message.content,
 					...(message.files && {
 						images: message.files
 							.filter((file) => file.type === 'image')
@@ -366,7 +468,7 @@
 								content: [
 									{
 										type: 'text',
-										text: message.content
+										text: message?.raContent ?? message.content
 									},
 									...message.files
 										.filter((file) => file.type === 'image')
@@ -378,7 +480,7 @@
 										}))
 								]
 						  }
-						: { content: message.content })
+						: { content: message?.raContent ?? message.content })
 				})),
 			seed: $settings?.options?.seed ?? undefined,
 			stop: $settings?.options?.stop ?? undefined,
@@ -494,73 +596,6 @@
 		}
 	};
 
-	const submitPrompt = async (userPrompt) => {
-		console.log('submitPrompt', $chatId);
-
-		if (selectedModels.includes('')) {
-			toast.error('Model not selected');
-		} else if (messages.length != 0 && messages.at(-1).done != true) {
-			// Response not done
-			console.log('wait');
-		} else {
-			// Reset chat message textarea height
-			document.getElementById('chat-textarea').style.height = '';
-
-			// Create user message
-			let userMessageId = uuidv4();
-			let userMessage = {
-				id: userMessageId,
-				parentId: messages.length !== 0 ? messages.at(-1).id : null,
-				childrenIds: [],
-				role: 'user',
-				content: userPrompt,
-				files: files.length > 0 ? files : undefined
-			};
-
-			// Add message to history and Set currentId to messageId
-			history.messages[userMessageId] = userMessage;
-			history.currentId = userMessageId;
-
-			// Append messageId to childrenIds of parent message
-			if (messages.length !== 0) {
-				history.messages[messages.at(-1).id].childrenIds.push(userMessageId);
-			}
-
-			// Wait until history/message have been updated
-			await tick();
-
-			// Create new chat if only one message in messages
-			if (messages.length == 1) {
-				if ($settings.saveChatHistory ?? true) {
-					chat = await createNewChat(localStorage.token, {
-						id: $chatId,
-						title: 'New Chat',
-						models: selectedModels,
-						system: $settings.system ?? undefined,
-						options: {
-							...($settings.options ?? {})
-						},
-						messages: messages,
-						history: history,
-						timestamp: Date.now()
-					});
-					await chats.set(await getChatList(localStorage.token));
-					await chatId.set(chat.id);
-				} else {
-					await chatId.set('local');
-				}
-				await tick();
-			}
-
-			// Reset chat input textarea
-			prompt = '';
-			files = [];
-
-			// Send prompt
-			await sendPrompt(userPrompt, userMessageId);
-		}
-	};
-
 	const stopResponse = () => {
 		stopResponseFlag = true;
 		console.log('stopResponse');