main.py 17.7 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

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
47
from apps.rag.utils import query_doc, query_collection, get_embedding_model_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_AUTO_UPDATE,
62
    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"])

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
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
80
81


82
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
83
app.state.RAG_EMBEDDING_MODEL_PATH = get_embedding_model_path(
Self Denial's avatar
Self Denial committed
84
85
    app.state.RAG_EMBEDDING_MODEL, RAG_EMBEDDING_MODEL_AUTO_UPDATE
)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
86
87


88
89
app.state.TOP_K = 4

90
91
app.state.sentence_transformer_ef = (
    embedding_functions.SentenceTransformerEmbeddingFunction(
92
        model_name=app.state.RAG_EMBEDDING_MODEL_PATH,
93
        device=DEVICE_TYPE,
94
95
    )
)
Timothy J. Baek's avatar
Timothy J. Baek committed
96

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

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

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


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


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

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

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


class EmbeddingModelUpdateForm(BaseModel):
    embedding_model: str


@app.post("/embedding/model/update")
async def update_embedding_model(
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
145
146
    status = True
    old_model_path = app.state.RAG_EMBEDDING_MODEL_PATH
147
    app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
148
149

    log.debug(f"form_data.embedding_model: {form_data.embedding_model}")
Self Denial's avatar
Self Denial committed
150
151
152
    log.info(
        f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
    )
153
154

    try:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
155
        app.state.RAG_EMBEDDING_MODEL_PATH = get_embedding_model_path(
Self Denial's avatar
Self Denial committed
156
157
            app.state.RAG_EMBEDDING_MODEL, True
        )
158
159
160
        app.state.sentence_transformer_ef = (
            embedding_functions.SentenceTransformerEmbeddingFunction(
                model_name=app.state.RAG_EMBEDDING_MODEL_PATH,
161
                device=DEVICE_TYPE,
162
            )
163
        )
Self Denial's avatar
Self Denial committed
164
    except Exception as e:
165
166
167
168
169
170
171
        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
172
        status = False
173

Self Denial's avatar
Self Denial committed
174
175
176
    log.debug(
        f"app.state.RAG_EMBEDDING_MODEL_PATH: {app.state.RAG_EMBEDDING_MODEL_PATH}"
    )
177
178
    log.debug(f"old_model_path: {old_model_path}")
    log.debug(f"status: {status}")
179
180

    return {
181
        "status": status,
182
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
183
        "embedding_model_path": app.state.RAG_EMBEDDING_MODEL_PATH,
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
    except Exception as e:
276
        log.exception(e)
277
278
279
280
        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
    except Exception as e:
320
        log.exception(e)
321
322
323
324
325
326
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


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

329
330
331
332
333
334
    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
335
336
337
338
339

    if len(docs) > 0:
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
340
341
342


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
343
    text, metadata, collection_name, overwrite: bool = False
344
345
346
347
348
349
) -> 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
350
    docs = text_splitter.create_documents([text], metadatas=[metadata])
351
352
353
354
    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
355

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

        return False


383
384
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
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
429
430
431
432
    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
433
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
434
435
436
437
438
439
    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)
440
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
441
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
442
443
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
444
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
445
446
        loader = UnstructuredEPubLoader(file_path)
    elif (
447
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
448
449
450
451
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
452
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
453
454
455
456
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
457
458
459
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
460
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
461
    else:
462
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
463
464
465
466
467
        known_type = False

    return loader, known_type


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

476
    log.info(f"file.content_type: {file.content_type}")
477
    try:
478
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
479
        filename = os.path.basename(unsanitized_filename)
480

Timothy J. Baek's avatar
Timothy J. Baek committed
481
        file_path = f"{UPLOAD_DIR}/{filename}"
482

483
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
484
        with open(file_path, "wb") as f:
485
486
487
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
488
489
490
491
492
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
493
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
494
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
495
496
497
498
499
500
501
502
503
504
505
506

        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
507
508
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
509
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
510
            )
511
    except Exception as e:
512
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
513
514
515
516
517
518
519
520
521
522
        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),
            )
523
524


525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
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
541
542
543
544
545
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
546
547
548
549
550
551
552
553
554
555

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


556
557
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
558
559
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
560
561
562
563
564
565
566
567
568
            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
569
570
571
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
572
573
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
                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,
                                                        )
598
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
599
600
601
602
603
604
605
606
607
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
608
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
609
                    pass
610

611
        except Exception as e:
612
            log.exception(e)
613
614
615
616

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
617
@app.get("/reset/db")
618
619
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
620
621
622


@app.get("/reset")
623
624
625
626
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
627
        try:
628
629
630
631
            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
632
        except Exception as e:
633
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
634

635
636
637
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
638
        log.exception(e)
639
640

    return True