open-webui/backend/apps/rag/main.py

271 lines
7.7 KiB
Python
Raw Normal View History

2024-01-07 08:40:51 +01:00
from fastapi import (
FastAPI,
Request,
Depends,
HTTPException,
status,
UploadFile,
File,
Form,
)
2024-01-07 07:07:20 +01:00
from fastapi.middleware.cors import CORSMiddleware
2024-01-07 10:40:36 +01:00
import os, shutil
2024-01-07 07:07:20 +01:00
2024-01-07 17:28:35 +01:00
# from chromadb.utils import embedding_functions
2024-01-07 07:07:20 +01:00
2024-01-07 18:05:52 +01:00
from langchain_community.document_loaders import (
WebBaseLoader,
TextLoader,
PyPDFLoader,
CSVLoader,
2024-01-07 22:56:01 +01:00
Docx2txtLoader,
2024-01-10 00:24:53 +01:00
UnstructuredWordDocumentLoader,
UnstructuredMarkdownLoader,
UnstructuredXMLLoader,
2024-01-07 18:05:52 +01:00
)
2024-01-07 07:59:22 +01:00
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
2024-01-07 18:33:34 +01:00
import time
2024-01-07 07:59:22 +01:00
2024-01-07 18:00:30 +01:00
from utils.misc import calculate_sha256
2024-01-07 11:46:12 +01:00
from utils.utils import get_current_user
2024-01-07 10:40:36 +01:00
from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
2024-01-07 07:59:22 +01:00
from constants import ERROR_MESSAGES
2024-01-07 17:28:35 +01:00
# EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
# model_name=EMBED_MODEL
# )
2024-01-07 07:07:20 +01:00
app = FastAPI()
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2024-01-07 08:40:51 +01:00
class CollectionNameForm(BaseModel):
2024-01-07 07:59:22 +01:00
collection_name: Optional[str] = "test"
2024-01-07 08:40:51 +01:00
class StoreWebForm(CollectionNameForm):
url: str
2024-01-07 10:40:36 +01:00
def store_data_in_vector_db(data, collection_name) -> bool:
2024-01-07 07:59:22 +01:00
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]
2024-01-07 10:40:36 +01:00
try:
2024-01-07 17:34:05 +01:00
collection = CHROMA_CLIENT.create_collection(name=collection_name)
2024-01-07 07:59:22 +01:00
2024-01-07 10:40:36 +01:00
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
2024-01-07 07:59:22 +01:00
2024-01-07 07:07:20 +01:00
@app.get("/")
async def get_status():
return {"status": True}
2024-01-07 07:59:22 +01:00
@app.get("/query/{collection_name}")
2024-01-07 11:46:12 +01:00
def query_collection(
collection_name: str,
query: str,
k: Optional[int] = 4,
user=Depends(get_current_user),
):
2024-01-07 10:59:00 +01:00
try:
collection = CHROMA_CLIENT.get_collection(
name=collection_name,
)
result = collection.query(query_texts=[query], n_results=k)
2024-01-07 07:59:22 +01:00
2024-01-07 10:59:00 +01:00
return result
except Exception as e:
print(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-01-07 07:59:22 +01:00
@app.post("/web")
2024-01-07 11:46:12 +01:00
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
2024-01-07 07:59:22 +01:00
# "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)
2024-01-08 10:26:15 +01:00
return {
"status": True,
"collection_name": form_data.collection_name,
"filename": form_data.url,
}
2024-01-07 07:59:22 +01:00
except Exception as e:
print(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
@app.post("/doc")
2024-01-07 11:46:12 +01:00
def store_doc(
2024-01-07 18:00:30 +01:00
collection_name: Optional[str] = Form(None),
2024-01-07 11:46:12 +01:00
file: UploadFile = File(...),
user=Depends(get_current_user),
):
2024-01-07 07:59:22 +01:00
# "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
2024-01-07 08:40:51 +01:00
2024-01-10 00:24:53 +01:00
print(file.content_type)
2024-01-07 22:56:01 +01:00
if file.content_type not in [
"application/pdf",
"text/plain",
"text/csv",
"text/xml",
"text/x-python",
2024-01-17 08:34:22 +01:00
"text/css",
2024-01-07 22:56:01 +01:00
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
2024-01-10 00:24:53 +01:00
"application/octet-stream",
2024-01-17 08:34:22 +01:00
"application/x-javascript",
2024-01-07 22:56:01 +01:00
]:
2024-01-07 08:40:51 +01:00
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
)
2024-01-17 08:34:22 +01:00
text_xml=["text/xml"]
octet_markdown=["md"]
octet_plain=[
"go", "py", "java", "sh", "bat", "ps1", "cmd", "js",
"css", "cpp", "hpp","h", "c", "cs", "sql", "log", "ini",
"pl" "pm", "r", "dart", "dockerfile", "env", "php", "hs",
"hsc", "lua", "nginxconf", "conf", "m", "mm", "plsql", "perl",
"rb", "rs", "db2", "scala", "bash", "swift", "vue"
]
file_ext=file.filename.split(".")[-1].lower()
if file.content_type == "application/octet-stream" and file_ext not in (octet_markdown + octet_plain):
2024-01-10 00:24:53 +01:00
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
)
2024-01-07 07:59:22 +01:00
try:
2024-01-07 08:40:51 +01:00
filename = file.filename
2024-01-07 10:40:36 +01:00
file_path = f"{UPLOAD_DIR}/{filename}"
2024-01-07 07:59:22 +01:00
contents = file.file.read()
2024-01-07 08:40:51 +01:00
with open(file_path, "wb") as f:
2024-01-07 07:59:22 +01:00
f.write(contents)
f.close()
2024-01-07 18:00:30 +01:00
f = open(file_path, "rb")
if collection_name == None:
collection_name = calculate_sha256(f)[:63]
f.close()
2024-01-07 08:40:51 +01:00
if file.content_type == "application/pdf":
loader = PyPDFLoader(file_path)
2024-01-07 22:56:01 +01:00
elif (
file.content_type
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
):
loader = Docx2txtLoader(file_path)
2024-01-07 18:05:52 +01:00
elif file.content_type == "text/csv":
loader = CSVLoader(file_path)
elif file.content_type in text_xml:
loader=UnstructuredXMLLoader(file_path)
elif file.content_type == "text/plain" or file.content_type.find("text/")>=0:
loader = TextLoader(file_path)
2024-01-10 00:24:53 +01:00
elif file.content_type == "application/octet-stream":
if file_ext in octet_markdown:
2024-01-10 00:24:53 +01:00
loader = UnstructuredMarkdownLoader(file_path)
if file_ext in octet_plain:
loader = TextLoader(file_path)
2024-01-17 08:34:22 +01:00
elif file.content_type == "application/x-javascript":
loader = TextLoader(file_path)
2024-01-07 08:40:51 +01:00
data = loader.load()
2024-01-07 10:40:36 +01:00
result = store_data_in_vector_db(data, collection_name)
if result:
2024-01-08 10:26:15 +01:00
return {
"status": True,
"collection_name": collection_name,
"filename": filename,
}
2024-01-07 10:40:36 +01:00
else:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=ERROR_MESSAGES.DEFAULT(),
)
2024-01-07 07:59:22 +01:00
except Exception as e:
print(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
2024-01-07 10:40:36 +01:00
@app.get("/reset/db")
2024-01-07 11:46:12 +01:00
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,
)
2024-01-07 10:40:36 +01:00
@app.get("/reset")
2024-01-07 18:15:45 +01:00
def reset(user=Depends(get_current_user)) -> bool:
2024-01-07 11:46:12 +01:00
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))
2024-01-07 10:40:36 +01:00
try:
2024-01-07 11:46:12 +01:00
CHROMA_CLIENT.reset()
2024-01-07 10:40:36 +01:00
except Exception as e:
2024-01-07 11:46:12 +01:00
print(e)
2024-01-07 10:40:36 +01:00
2024-01-07 18:15:45 +01:00
return True
2024-01-07 11:46:12 +01:00
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)