main.py 14.6 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
    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
#    sentence_transformer_ef = SentenceTransformer(
#        model_name_or_path=RAG_EMBEDDING_MODEL,
#        cache_folder=RAG_EMBEDDING_MODEL_DIR,
Jannik Streidl's avatar
Jannik Streidl committed
74
#        device=DEVICE_TYPE,
75
#    )
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
app.state.PDF_EXTRACT_IMAGES = False
Timothy J. Baek's avatar
Timothy J. Baek committed
81
82
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
83
app.state.RAG_TEMPLATE = RAG_TEMPLATE
84
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
85
86
app.state.TOP_K = 4

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

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

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

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


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


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


114
def store_data_in_vector_db(data, collection_name, overwrite: bool = False) -> bool:
115
    text_splitter = RecursiveCharacterTextSplitter(
Timothy J. Baek's avatar
Timothy J. Baek committed
116
        chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP
117
118
119
120
121
122
    )
    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
123
    try:
124
125
126
127
128
        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)
129

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

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

        return False
145
146


Timothy J. Baek's avatar
Timothy J. Baek committed
147
148
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
149
150
151
152
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
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
        "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,
178
            device=DEVICE_TYPE,
179
180
181
182
183
        )
    )
    return {
        "status": True,
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
186
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
187
188
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
191
192
193
194
195
        "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
196
197
198
199
200
201
202
203
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
204
205
206
207
208
209
210
211
212
213
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
214
215
216

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
217
218
219
220
221
        "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
222
    }
223
224


Timothy J. Baek's avatar
Timothy J. Baek committed
225
226
227
228
229
230
231
232
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


233
234
235
236
237
238
239
@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
240
241


242
243
244
245
246
247
248
249
250
251
252
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
253
    return {"status": True, "template": app.state.RAG_TEMPLATE}
254
255


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


262
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
263
def query_doc_handler(
264
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
267

268
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
269
270
271
272
        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,
273
274
            embedding_function=app.state.sentence_transformer_ef,
        )
275
276
277
278
279
280
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
281
282


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
283
284
285
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
286
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
287
288


289
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
290
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
291
292
293
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
296
297
298
    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,
299
    )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
300
301


302
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
303
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
304
305
306
307
    # "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
308
309
310
311
312

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

313
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
314
315
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
316
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
317
318
            "filename": form_data.url,
        }
319
320
321
322
323
324
325
326
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


327
328
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
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
367
368
369
370
371
372
373
374
375
376
    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
377
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
378
379
380
381
382
383
384
385
    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)
386
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
387
388
        loader = UnstructuredEPubLoader(file_path)
    elif (
389
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
390
391
392
393
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
394
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
395
396
397
398
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
399
400
401
    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
402
403
404
405
406
407
408
409
        loader = TextLoader(file_path)
    else:
        loader = TextLoader(file_path)
        known_type = False

    return loader, known_type


410
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
411
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
412
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
413
414
415
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
416
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
417

Timothy J. Baek's avatar
Timothy J. Baek committed
418
    print(file.content_type)
419
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
420
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
421
        file_path = f"{UPLOAD_DIR}/{filename}"
422
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
423
        with open(file_path, "wb") as f:
424
425
426
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
427
428
429
430
431
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

432
        loader, known_type = get_loader(file.filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
433
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
434
435
436
        result = store_data_in_vector_db(data, collection_name)

        if result:
Timothy J. Baek's avatar
Timothy J. Baek committed
437
438
439
440
            return {
                "status": True,
                "collection_name": collection_name,
                "filename": filename,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
441
                "known_type": known_type,
Timothy J. Baek's avatar
Timothy J. Baek committed
442
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
443
444
445
446
447
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
448
449
    except Exception as e:
        print(e)
Dave Bauman's avatar
Dave Bauman committed
450
451
452
453
454
455
456
457
458
459
        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),
            )
460
461


462
463
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
464
465
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
466
467
468
469
470
471
472
473
474
            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
475
476
477
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
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
503
504
505
506
507
508
509
510
511
512
                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 "{}"
                                    ),
                                }
                            ),
                        )

513
514
        except Exception as e:
            print(e)
515
516
517
518

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
519
@app.get("/reset/db")
520
521
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
522
523
524


@app.get("/reset")
525
526
527
528
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
529
        try:
530
531
532
533
            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
534
        except Exception as e:
535
            print("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
536

537
538
539
540
541
542
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
        print(e)

    return True