main.py 16.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
11
import os, shutil, logging
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,
24
    BSHTMLLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
25
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
26
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
27
28
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
29
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
30
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
31
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
32
)
33
34
35
36
from langchain.text_splitter import RecursiveCharacterTextSplitter

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

41

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

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

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

69
70
from constants import ERROR_MESSAGES

71
72
73
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

74
#
75
# if RAG_EMBEDDING_MODEL:
76
77
78
79
80
#    sentence_transformer_ef = SentenceTransformer(
#        model_name_or_path=RAG_EMBEDDING_MODEL,
#        cache_folder=RAG_EMBEDDING_MODEL_DIR,
#        device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
#    )
81

Timothy J. Baek's avatar
Timothy J. Baek committed
82
83
84

app = FastAPI()

Timothy J. Baek's avatar
Timothy J. Baek committed
85
app.state.PDF_EXTRACT_IMAGES = False
Timothy J. Baek's avatar
Timothy J. Baek committed
86
87
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
88
app.state.RAG_TEMPLATE = RAG_TEMPLATE
89
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
90
91
app.state.TOP_K = 4

92
93
94
95
96
97
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
98

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

Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102
103
104
105
106
107
108
109
110
origins = ["*"]

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


Timothy J. Baek's avatar
Timothy J. Baek committed
111
class CollectionNameForm(BaseModel):
112
113
114
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
115
116
117
class StoreWebForm(CollectionNameForm):
    url: str

Timothy J. Baek's avatar
Timothy J. Baek committed
118
119
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
120
121
122
123
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        "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
156
157
158
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
159
160
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
161
162
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
163
164
165
166
167
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
170
171
172
173
174
175
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
176
177
178
179
180
181
182
183
184
185
class ConfigUpdateForm(BaseModel):
    pdf_extract_images: bool
    chunk: ChunkParamUpdateForm


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
    app.state.PDF_EXTRACT_IMAGES = form_data.pdf_extract_images
    app.state.CHUNK_SIZE = form_data.chunk.chunk_size
    app.state.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
188

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
191
192
193
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
194
    }
195
196


Timothy J. Baek's avatar
Timothy J. Baek committed
197
198
199
200
201
202
203
204
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


205
206
207
208
209
210
211
@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
212
213


214
215
216
217
218
219
220
221
222
223
224
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
225
    return {"status": True, "template": app.state.RAG_TEMPLATE}
226
227


228
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
229
230
    collection_name: str
    query: str
231
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
232
233


234
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
235
def query_doc_handler(
236
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
237
238
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
239

240
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
241
242
243
244
        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,
245
246
            embedding_function=app.state.sentence_transformer_ef,
        )
247
    except Exception as e:
248
        log.exception(e)
249
250
251
252
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
253
254


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
255
256
257
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
258
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
259
260


261
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
262
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
263
264
265
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
266
267
268
269
270
    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,
271
    )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
272
273


274
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
275
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
276
277
278
279
    # "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
280
281
282
283
284

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

285
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
286
287
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
288
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
289
290
            "filename": form_data.url,
        }
291
    except Exception as e:
292
        log.exception(e)
293
294
295
296
297
298
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


299
300
301
302
303
304
305
306
307
308
309
def store_data_in_vector_db(data, collection_name, overwrite: bool = False) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
    docs = text_splitter.split_documents(data)
    return store_docs_in_vector_db(docs, collection_name, overwrite)


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
310
    text, metadata, collection_name, overwrite: bool = False
311
312
313
314
315
316
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
317
    docs = text_splitter.create_documents([text], metadatas=[metadata])
318
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
    return store_docs_in_vector_db(docs, collection_name, overwrite)


def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> bool:
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]

    try:
        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)

        collection = CHROMA_CLIENT.create_collection(
            name=collection_name,
            embedding_function=app.state.sentence_transformer_ef,
        )

        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


349
350
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
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
    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
399
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
400
401
402
403
404
405
    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)
406
    elif file_ext in ["htm", "html"]:
407
        loader = BSHTMLLoader(file_path)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
408
409
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
410
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
411
412
        loader = UnstructuredEPubLoader(file_path)
    elif (
413
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
414
415
416
417
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
418
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
419
420
421
422
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
423
424
425
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
426
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
427
    else:
428
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
429
430
431
432
433
        known_type = False

    return loader, known_type


434
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
435
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
436
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
437
438
439
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
440
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
441

442
    log.info(f"file.content_type: {file.content_type}")
443
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
444
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
445
        file_path = f"{UPLOAD_DIR}/{filename}"
446
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
447
        with open(file_path, "wb") as f:
448
449
450
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
451
452
453
454
455
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

456
        loader, known_type = get_loader(file.filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
457
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
458
459
460
        result = store_data_in_vector_db(data, collection_name)

        if result:
Timothy J. Baek's avatar
Timothy J. Baek committed
461
462
463
464
            return {
                "status": True,
                "collection_name": collection_name,
                "filename": filename,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
465
                "known_type": known_type,
Timothy J. Baek's avatar
Timothy J. Baek committed
466
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
467
468
469
470
471
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
472
    except Exception as e:
473
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
474
475
476
477
478
479
480
481
482
483
        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),
            )
484
485


486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
class TextRAGForm(BaseModel):
    name: str
    content: str
    collection_name: Optional[str] = None


@app.post("/text")
def store_text(
    form_data: TextRAGForm,
    user=Depends(get_current_user),
):

    collection_name = form_data.collection_name
    if collection_name == None:
        collection_name = calculate_sha256_string(form_data.content)

Timothy J. Baek's avatar
Timothy J. Baek committed
502
503
504
505
506
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
507
508
509
510
511
512
513
514
515
516

    if result:
        return {"status": True, "collection_name": collection_name}
    else:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=ERROR_MESSAGES.DEFAULT(),
        )


517
518
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
519
520
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
521
522
523
524
525
526
527
528
529
            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
530
531
532
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
                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 "{}"
                                    ),
                                }
                            ),
                        )

568
        except Exception as e:
569
            log.exception(e)
570
571
572
573

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
574
@app.get("/reset/db")
575
576
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
577
578
579


@app.get("/reset")
580
581
582
583
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
584
        try:
585
586
587
588
            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
589
        except Exception as e:
590
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
591

592
593
594
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
595
        log.exception(e)
596
597

    return True