main.py 7.65 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
4
5
6
7
8
9
10
from fastapi import (
    FastAPI,
    Request,
    Depends,
    HTTPException,
    status,
    UploadFile,
    File,
    Form,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
11
from fastapi.middleware.cors import CORSMiddleware
Timothy J. Baek's avatar
Timothy J. Baek committed
12
import os, shutil
Timothy J. Baek's avatar
Timothy J. Baek committed
13

Timothy J. Baek's avatar
Timothy J. Baek committed
14
# from chromadb.utils import embedding_functions
Timothy J. Baek's avatar
Timothy J. Baek committed
15

Timothy J. Baek's avatar
Timothy J. Baek committed
16
17
18
19
20
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
21
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
22
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
23
24
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
25
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
26
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
27
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
28
)
29
30
31
32
33
34
35
36
37
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
38
import time
39

Timothy J. Baek's avatar
Timothy J. Baek committed
40
from utils.misc import calculate_sha256
Timothy J. Baek's avatar
Timothy J. Baek committed
41
from utils.utils import get_current_user
Timothy J. Baek's avatar
Timothy J. Baek committed
42
from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
43
44
from constants import ERROR_MESSAGES

Timothy J. Baek's avatar
Timothy J. Baek committed
45
46
47
# EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
#     model_name=EMBED_MODEL
# )
Timothy J. Baek's avatar
Timothy J. Baek committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61

app = FastAPI()

origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
62
class CollectionNameForm(BaseModel):
63
64
65
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
66
67
68
69
class StoreWebForm(CollectionNameForm):
    url: str


Timothy J. Baek's avatar
Timothy J. Baek committed
70
def store_data_in_vector_db(data, collection_name) -> bool:
71
72
73
74
75
76
77
78
    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]

Timothy J. Baek's avatar
Timothy J. Baek committed
79
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
80
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
81

Timothy J. Baek's avatar
Timothy J. Baek committed
82
83
84
85
86
87
88
89
90
91
        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
92
93


Timothy J. Baek's avatar
Timothy J. Baek committed
94
95
96
@app.get("/")
async def get_status():
    return {"status": True}
97
98
99


@app.get("/query/{collection_name}")
Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102
103
104
105
def query_collection(
    collection_name: str,
    query: str,
    k: Optional[int] = 4,
    user=Depends(get_current_user),
):
106
107
108
109
110
    try:
        collection = CHROMA_CLIENT.get_collection(
            name=collection_name,
        )
        result = collection.query(query_texts=[query], n_results=k)
111

112
113
114
115
116
117
118
        return result
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
119
120
121


@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
122
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
123
124
125
126
127
    # "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)
Timothy J. Baek's avatar
Timothy J. Baek committed
128
129
130
131
132
        return {
            "status": True,
            "collection_name": form_data.collection_name,
            "filename": form_data.url,
        }
133
134
135
136
137
138
139
140
141
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
142
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
143
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
144
145
146
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
147
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
148

Timothy J. Baek's avatar
Timothy J. Baek committed
149
    print(file.content_type)
Marclass's avatar
Marclass committed
150
151
    
    text_xml=["xml"]
152
    octet_markdown=["md"]
Marclass's avatar
Marclass committed
153
    known_source_ext=[
154
        "go", "py", "java", "sh", "bat", "ps1", "cmd", "js", "ts",
155
        "css", "cpp", "hpp","h", "c", "cs", "sql", "log", "ini",
156
        "pl", "pm", "r", "dart", "dockerfile", "env", "php", "hs",
157
        "hsc", "lua", "nginxconf", "conf", "m", "mm", "plsql", "perl",
Marclass's avatar
Marclass committed
158
        "rb", "rs", "db2", "scala", "bash", "swift", "vue", "svelte"
159
        ]
Marclass's avatar
Marclass committed
160
161
    docx_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    known_doc_ext=["doc","docx"]
Marclass's avatar
Marclass committed
162
163
164
    excel_types=["application/vnd.ms-excel", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"]
    known_excel_ext=["xls", "xlsx"]

165
    file_ext=file.filename.split(".")[-1].lower()
Marclass's avatar
Marclass committed
166
167
    known_type=True
    
168
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
169
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
170
        file_path = f"{UPLOAD_DIR}/{filename}"
171
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
172
        with open(file_path, "wb") as f:
173
174
175
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
176
177
178
179
180
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

Marclass's avatar
Marclass committed
181
        if file_ext=="pdf":
Timothy J. Baek's avatar
Timothy J. Baek committed
182
            loader = PyPDFLoader(file_path)
Marclass's avatar
Marclass committed
183
        elif (file.content_type ==docx_type or file_ext in known_doc_ext):
Timothy J. Baek's avatar
Timothy J. Baek committed
184
            loader = Docx2txtLoader(file_path)
Marclass's avatar
Marclass committed
185
        elif file_ext=="csv":
Timothy J. Baek's avatar
Timothy J. Baek committed
186
            loader = CSVLoader(file_path)
Marclass's avatar
Marclass committed
187
188
        elif (file.content_type in excel_types or file_ext in known_excel_ext):
            loader = UnstructuredExcelLoader(file_path)
Marclass's avatar
Marclass committed
189
190
        elif file_ext=="rst":
            loader = UnstructuredRSTLoader(file_path, mode="elements")
Marclass's avatar
Marclass committed
191
        elif file_ext in text_xml:
192
            loader=UnstructuredXMLLoader(file_path)
Marclass's avatar
Marclass committed
193
        elif file_ext in known_source_ext or file.content_type.find("text/")>=0:
194
            loader = TextLoader(file_path)
Marclass's avatar
Marclass committed
195
196
        elif file_ext in octet_markdown:
            loader = UnstructuredMarkdownLoader(file_path)
Dave Bauman's avatar
Dave Bauman committed
197
198
        elif file.content_type == "application/epub+zip":
            loader = UnstructuredEPubLoader(file_path)
Marclass's avatar
Marclass committed
199
        else:
Marclass's avatar
Marclass committed
200
            loader = TextLoader(file_path)
Marclass's avatar
Marclass committed
201
202
            known_type=False

Timothy J. Baek's avatar
Timothy J. Baek committed
203
204

        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
        result = store_data_in_vector_db(data, collection_name)

        if result:
Timothy J. Baek's avatar
Timothy J. Baek committed
208
209
210
211
            return {
                "status": True,
                "collection_name": collection_name,
                "filename": filename,
Marclass's avatar
Marclass committed
212
                "known_type":known_type,
Timothy J. Baek's avatar
Timothy J. Baek committed
213
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
214
215
216
217
218
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
219
220
    except Exception as e:
        print(e)
Dave Bauman's avatar
Dave Bauman committed
221
222
223
224
225
226
227
228
229
230
        if "No pandoc was found" in str(e):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.PANDOC_NOT_INSTALLED,
            )
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e),
            )
231
232


Timothy J. Baek's avatar
Timothy J. Baek committed
233
@app.get("/reset/db")
Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
236
237
238
239
240
241
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,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
242
243
244


@app.get("/reset")
245
def reset(user=Depends(get_current_user)) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
246
247
248
249
250
251
252
253
254
255
256
257
    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))

Timothy J. Baek's avatar
Timothy J. Baek committed
258
        try:
Timothy J. Baek's avatar
Timothy J. Baek committed
259
            CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
260
        except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
261
            print(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
262

263
        return True
Timothy J. Baek's avatar
Timothy J. Baek committed
264
265
266
267
268
    else:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )