main.py 18.4 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, re
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 chromadb.utils import embedding_functions
Timothy J. Baek's avatar
Timothy J. Baek committed
17

Timothy J. Baek's avatar
Timothy J. Baek committed
18
19
20
21
22
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
23
    BSHTMLLoader,
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

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

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
from config import (
57
    SRC_LOG_LEVELS,
58
59
    UPLOAD_DIR,
    DOCS_DIR,
60
    RAG_EMBEDDING_MODEL,
61
    RAG_EMBEDDING_MODEL_DEVICE_TYPE,
62
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
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"])

Timothy J. Baek's avatar
Timothy J. Baek committed
74
75
app = FastAPI()

Timothy J. Baek's avatar
Timothy J. Baek committed
76
app.state.PDF_EXTRACT_IMAGES = False
Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
79
app.state.RAG_TEMPLATE = RAG_TEMPLATE
80
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
Self Denial's avatar
Self Denial committed
81
82
83
app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path(
    app.state.RAG_EMBEDDING_MODEL, RAG_EMBEDDING_MODEL_AUTO_UPDATE
)
84
85
app.state.TOP_K = 4

86
87
app.state.sentence_transformer_ef = (
    embedding_functions.SentenceTransformerEmbeddingFunction(
88
        model_name=app.state.RAG_EMBEDDING_MODEL_PATH,
89
90
91
        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
class StoreWebForm(CollectionNameForm):
    url: str

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

Timothy J. Baek's avatar
Timothy J. Baek committed
113
114
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
115
116
117
118
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
119
120
121
122
123
124
125
126
127
128
        "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,
129
        "embedding_model_path": app.state.RAG_EMBEDDING_MODEL_PATH,
130
131
132
133
134
135
136
137
138
139
140
    }


class EmbeddingModelUpdateForm(BaseModel):
    embedding_model: str


@app.post("/embedding/model/update")
async def update_embedding_model(
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
141
142
    status = True
    old_model_path = app.state.RAG_EMBEDDING_MODEL_PATH
143
    app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
144
145

    log.debug(f"form_data.embedding_model: {form_data.embedding_model}")
Self Denial's avatar
Self Denial committed
146
147
148
    log.info(
        f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
    )
149
150

    try:
Self Denial's avatar
Self Denial committed
151
152
153
        app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path(
            app.state.RAG_EMBEDDING_MODEL, True
        )
154
155
156
157
158
        app.state.sentence_transformer_ef = (
            embedding_functions.SentenceTransformerEmbeddingFunction(
                model_name=app.state.RAG_EMBEDDING_MODEL_PATH,
                device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
            )
159
        )
Self Denial's avatar
Self Denial committed
160
    except Exception as e:
161
162
163
164
165
166
167
        log.exception(f"Problem updating embedding model: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=e,
        )

    if app.state.RAG_EMBEDDING_MODEL_PATH == old_model_path:
Self Denial's avatar
Self Denial committed
168
        status = False
169

Self Denial's avatar
Self Denial committed
170
171
172
    log.debug(
        f"app.state.RAG_EMBEDDING_MODEL_PATH: {app.state.RAG_EMBEDDING_MODEL_PATH}"
    )
173
174
    log.debug(f"old_model_path: {old_model_path}")
    log.debug(f"status: {status}")
175
176

    return {
177
        "status": status,
178
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
179
        "embedding_model_path": app.state.RAG_EMBEDDING_MODEL_PATH,
Timothy J. Baek's avatar
Timothy J. Baek committed
180
181
182
    }


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


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


Timothy J. Baek's avatar
Timothy J. Baek committed
200
201
202
203
204
205
206
207
208
209
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
210
211
212

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
213
214
215
216
217
        "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
218
    }
219
220


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


229
230
231
232
233
234
235
@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
236
237


238
239
240
241
242
243
244
245
246
247
248
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
249
    return {"status": True, "template": app.state.RAG_TEMPLATE}
250
251


252
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
253
254
    collection_name: str
    query: str
255
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
256
257


258
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
259
def query_doc_handler(
260
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
261
262
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
263

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


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
279
280
281
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
282
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
283
284


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


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

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

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


323
def store_data_in_vector_db(data, collection_name, overwrite: bool = False) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
324

325
326
327
328
329
330
    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)
Timothy J. Baek's avatar
Timothy J. Baek committed
331
332
333
334
335

    if len(docs) > 0:
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
336
337
338


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
339
    text, metadata, collection_name, overwrite: bool = False
340
341
342
343
344
345
) -> 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
346
    docs = text_splitter.create_documents([text], metadatas=[metadata])
347
348
349
350
    return store_docs_in_vector_db(docs, collection_name, overwrite)


def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
351

352
353
354
355
356
357
358
    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:
359
                    log.info(f"deleting existing collection {collection_name}")
360
361
362
363
364
365
366
367
368
369
370
371
                    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:
372
        log.exception(e)
373
374
375
376
377
378
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


379
380
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
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
    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
429
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
430
431
432
433
434
435
    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)
436
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
437
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
438
439
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
440
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
441
442
        loader = UnstructuredEPubLoader(file_path)
    elif (
443
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
444
445
446
447
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
448
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
449
450
451
452
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
453
454
455
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
456
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
457
    else:
458
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
459
460
461
462
463
        known_type = False

    return loader, known_type


464
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
465
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
466
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
467
468
469
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
470
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
471

472
    log.info(f"file.content_type: {file.content_type}")
473
    try:
474
475
        is_valid_filename = True
        unsanitized_filename = file.filename
476
        if re.search(r'[\\/:"\*\?<>|\n\t ]', unsanitized_filename) is not None:
477
478
            is_valid_filename = False

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
479
        unvalidated_file_path = f"{UPLOAD_DIR}/{unsanitized_filename}"
480
481
482
483
484
485
486
487
488
489
490
491
        dereferenced_file_path = str(Path(unvalidated_file_path).resolve(strict=False))
        if not dereferenced_file_path.startswith(UPLOAD_DIR):
            is_valid_filename = False

        if is_valid_filename:
            file_path = dereferenced_file_path
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(),
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
492
        filename = file.filename
493
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
494
        with open(file_path, "wb") as f:
495
496
497
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
498
499
500
501
502
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

503
        loader, known_type = get_loader(file.filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
504
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
505
506
507
508
509
510
511
512
513
514
515
516

        try:
            result = store_data_in_vector_db(data, collection_name)

            if result:
                return {
                    "status": True,
                    "collection_name": collection_name,
                    "filename": filename,
                    "known_type": known_type,
                }
        except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
517
518
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
519
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
520
            )
521
    except Exception as e:
522
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
523
524
525
526
527
528
529
530
531
532
        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),
            )
533
534


535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
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
551
552
553
554
555
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
556
557
558
559
560
561
562
563
564
565

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


566
567
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
568
569
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
570
571
572
573
574
575
576
577
578
            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
579
580
581
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
582
583
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
                try:
                    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,
                                                        )
608
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
609
610
611
612
613
614
615
616
617
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
618
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
619
                    pass
620

621
        except Exception as e:
622
            log.exception(e)
623
624
625
626

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
627
@app.get("/reset/db")
628
629
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
630
631
632


@app.get("/reset")
633
634
635
636
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
637
        try:
638
639
640
641
            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
642
        except Exception as e:
643
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
644

645
646
647
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
648
        log.exception(e)
649
650

    return True