Merge pull request #333 from ollama-webui/rag

feat: RAG support

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",
-                                     "http://localhost:11434/api")
+OLLAMA_API_BASE_URL = os.environ.get(
+    "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),
-          name="spa-static-files")
+
+app.mount("/", SPAStaticFiles(directory="../build", html=True), name="spa-static-files")

backend/requirements.txt

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