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

from pathlib import Path
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
14
from typing import List
Timothy J. Baek's avatar
Timothy J. Baek committed
15

16
from sentence_transformers import SentenceTransformer
17
from chromadb.utils import embedding_functions
Timothy J. Baek's avatar
Timothy J. Baek committed
18

Timothy J. Baek's avatar
Timothy J. Baek committed
19
20
21
22
23
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
24
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
25
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
26
27
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
28
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
29
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
30
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
31
)
32
33
34
35
from langchain.text_splitter import RecursiveCharacterTextSplitter

from pydantic import BaseModel
from typing import Optional
36
import mimetypes
37
import uuid
38
39
import json

40

41
42
43
44
45
from apps.web.models.documents import (
    Documents,
    DocumentForm,
    DocumentResponse,
)
Jannik Streidl's avatar
Jannik Streidl committed
46

Timothy J. Baek's avatar
Timothy J. Baek committed
47
48
from apps.rag.utils import query_doc, query_collection

49
50
51
52
53
54
from utils.misc import (
    calculate_sha256,
    calculate_sha256_string,
    sanitize_filename,
    extract_folders_after_data_docs,
)
55
from utils.utils import get_current_user, get_admin_user
56
57
58
from config import (
    UPLOAD_DIR,
    DOCS_DIR,
59
    RAG_EMBEDDING_MODEL,
60
    RAG_EMBEDDING_MODEL_DEVICE_TYPE,
61
62
63
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
64
    RAG_TEMPLATE,
65
)
66

67
68
from constants import ERROR_MESSAGES

69
#
70
# if RAG_EMBEDDING_MODEL:
71
72
73
74
75
#    sentence_transformer_ef = SentenceTransformer(
#        model_name_or_path=RAG_EMBEDDING_MODEL,
#        cache_folder=RAG_EMBEDDING_MODEL_DIR,
#        device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
#    )
76

Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
79

app = FastAPI()

Timothy J. Baek's avatar
Timothy J. Baek committed
80
81
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
82
app.state.RAG_TEMPLATE = RAG_TEMPLATE
83
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
84
85
app.state.TOP_K = 4

86
87
88
89
90
91
app.state.sentence_transformer_ef = (
    embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name=app.state.RAG_EMBEDDING_MODEL,
        device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
    )
)
Timothy J. Baek's avatar
Timothy J. Baek committed
92

Timothy J. Baek's avatar
Timothy J. Baek committed
93

Timothy J. Baek's avatar
Timothy J. Baek committed
94
95
96
97
98
99
100
101
102
103
104
origins = ["*"]

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


Timothy J. Baek's avatar
Timothy J. Baek committed
105
class CollectionNameForm(BaseModel):
106
107
108
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
109
110
111
112
class StoreWebForm(CollectionNameForm):
    url: str


113
def store_data_in_vector_db(data, collection_name, overwrite: bool = False) -> bool:
114
    text_splitter = RecursiveCharacterTextSplitter(
Timothy J. Baek's avatar
Timothy J. Baek committed
115
        chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP
116
117
118
119
120
121
    )
    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
122
    try:
123
124
125
126
127
        if overwrite:
            for collection in CHROMA_CLIENT.list_collections():
                if collection_name == collection.name:
                    print(f"deleting existing collection {collection_name}")
                    CHROMA_CLIENT.delete_collection(name=collection_name)
128

129
130
131
132
        collection = CHROMA_CLIENT.create_collection(
            name=collection_name,
            embedding_function=app.state.sentence_transformer_ef,
        )
133

Timothy J. Baek's avatar
Timothy J. Baek committed
134
        collection.add(
135
136
            documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts]
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
137
138
139
140
141
142
143
        return True
    except Exception as e:
        print(e)
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False
144
145


Timothy J. Baek's avatar
Timothy J. Baek committed
146
147
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
148
149
150
151
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
        "template": app.state.RAG_TEMPLATE,
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
    }


@app.get("/embedding/model")
async def get_embedding_model(user=Depends(get_admin_user)):
    return {
        "status": True,
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
    }


class EmbeddingModelUpdateForm(BaseModel):
    embedding_model: str


@app.post("/embedding/model/update")
async def update_embedding_model(
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
    app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
    app.state.sentence_transformer_ef = (
        embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=app.state.RAG_EMBEDDING_MODEL,
            device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
        )
    )

    return {
        "status": True,
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
    }


@app.get("/chunk")
async def get_chunk_params(user=Depends(get_admin_user)):
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
    }


class ChunkParamUpdateForm(BaseModel):
    chunk_size: int
    chunk_overlap: int


@app.post("/chunk/update")
async def update_chunk_params(
    form_data: ChunkParamUpdateForm, user=Depends(get_admin_user)
):
    app.state.CHUNK_SIZE = form_data.chunk_size
    app.state.CHUNK_OVERLAP = form_data.chunk_overlap

    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
    }
213
214


Timothy J. Baek's avatar
Timothy J. Baek committed
215
216
217
218
219
220
221
222
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


223
224
225
226
227
228
229
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
        "k": app.state.TOP_K,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
230
231


232
233
234
235
236
237
238
239
240
241
242
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
    template: Optional[str] = None


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
    app.state.RAG_TEMPLATE = form_data.template if form_data.template else RAG_TEMPLATE
    app.state.TOP_K = form_data.k if form_data.k else 4
Timothy J. Baek's avatar
Timothy J. Baek committed
243
    return {"status": True, "template": app.state.RAG_TEMPLATE}
244
245


246
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
247
248
    collection_name: str
    query: str
249
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
250
251


252
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
253
def query_doc_handler(
254
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
255
256
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
257

258
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
259
260
261
262
        return query_doc(
            collection_name=form_data.collection_name,
            query=form_data.query,
            k=form_data.k if form_data.k else app.state.TOP_K,
263
264
            embedding_function=app.state.sentence_transformer_ef,
        )
265
266
267
268
269
270
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
271
272


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
273
274
275
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
276
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
277
278


279
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
280
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
281
282
283
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
284
285
286
287
288
    return query_collection(
        collection_names=form_data.collection_names,
        query=form_data.query,
        k=form_data.k if form_data.k else app.state.TOP_K,
        embedding_function=app.state.sentence_transformer_ef,
289
    )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
290
291


292
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
293
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
294
295
296
297
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
        loader = WebBaseLoader(form_data.url)
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
298
299
300
301
302

        collection_name = form_data.collection_name
        if collection_name == "":
            collection_name = calculate_sha256_string(form_data.url)[:63]

303
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
304
305
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
306
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
307
308
            "filename": form_data.url,
        }
309
310
311
312
313
314
315
316
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


317
318
def get_loader(filename: str, file_content_type: str, file_path: str):
    file_ext = filename.split(".")[-1].lower()
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
    known_type = True

    known_source_ext = [
        "go",
        "py",
        "java",
        "sh",
        "bat",
        "ps1",
        "cmd",
        "js",
        "ts",
        "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",
        "svelte",
    ]

    if file_ext == "pdf":
Timothy J. Baek's avatar
Timothy J. Baek committed
367
        loader = PyPDFLoader(file_path, extract_images=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
368
369
370
371
372
373
374
375
    elif file_ext == "csv":
        loader = CSVLoader(file_path)
    elif file_ext == "rst":
        loader = UnstructuredRSTLoader(file_path, mode="elements")
    elif file_ext == "xml":
        loader = UnstructuredXMLLoader(file_path)
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
376
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
377
378
        loader = UnstructuredEPubLoader(file_path)
    elif (
379
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
380
381
382
383
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
384
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
385
386
387
388
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
389
390
391
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
392
393
394
395
396
397
398
399
        loader = TextLoader(file_path)
    else:
        loader = TextLoader(file_path)
        known_type = False

    return loader, known_type


400
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
401
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
402
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
403
404
405
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
406
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
407

Timothy J. Baek's avatar
Timothy J. Baek committed
408
    print(file.content_type)
409
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
410
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
411
        file_path = f"{UPLOAD_DIR}/{filename}"
412
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
413
        with open(file_path, "wb") as f:
414
415
416
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
417
418
419
420
421
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

422
        loader, known_type = get_loader(file.filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
423
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
424
425
426
        result = store_data_in_vector_db(data, collection_name)

        if result:
Timothy J. Baek's avatar
Timothy J. Baek committed
427
428
429
430
            return {
                "status": True,
                "collection_name": collection_name,
                "filename": filename,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
431
                "known_type": known_type,
Timothy J. Baek's avatar
Timothy J. Baek committed
432
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
433
434
435
436
437
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
438
439
    except Exception as e:
        print(e)
Dave Bauman's avatar
Dave Bauman committed
440
441
442
443
444
445
446
447
448
449
        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),
            )
450
451


452
453
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
454
455
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
456
457
458
459
460
461
462
463
464
            if path.is_file() and not path.name.startswith("."):
                tags = extract_folders_after_data_docs(path)
                filename = path.name
                file_content_type = mimetypes.guess_type(path)

                f = open(path, "rb")
                collection_name = calculate_sha256(f)[:63]
                f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
465
466
467
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
                data = loader.load()

                result = store_data_in_vector_db(data, collection_name)

                if result:
                    sanitized_filename = sanitize_filename(filename)
                    doc = Documents.get_doc_by_name(sanitized_filename)

                    if doc == None:
                        doc = Documents.insert_new_doc(
                            user.id,
                            DocumentForm(
                                **{
                                    "name": sanitized_filename,
                                    "title": filename,
                                    "collection_name": collection_name,
                                    "filename": filename,
                                    "content": (
                                        json.dumps(
                                            {
                                                "tags": list(
                                                    map(
                                                        lambda name: {"name": name},
                                                        tags,
                                                    )
                                                )
                                            }
                                        )
                                        if len(tags)
                                        else "{}"
                                    ),
                                }
                            ),
                        )

503
504
        except Exception as e:
            print(e)
505
506
507
508

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
509
@app.get("/reset/db")
510
511
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
512
513
514


@app.get("/reset")
515
516
517
518
def reset(user=Depends(get_admin_user)) -> bool:
    folder = f"{UPLOAD_DIR}"
    for filename in os.listdir(folder):
        file_path = os.path.join(folder, filename)
Timothy J. Baek's avatar
Timothy J. Baek committed
519
        try:
520
521
522
523
            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)
Timothy J. Baek's avatar
Timothy J. Baek committed
524
        except Exception as e:
525
            print("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
526

527
528
529
530
531
532
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
        print(e)

    return True