main.py 18.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, 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
81
app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path(app.state.RAG_EMBEDDING_MODEL, RAG_EMBEDDING_MODEL_AUTO_UPDATE)
82
83
app.state.TOP_K = 4

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

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

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

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


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


Timothy J. Baek's avatar
Timothy J. Baek committed
107
108
109
class StoreWebForm(CollectionNameForm):
    url: str

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

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


class EmbeddingModelUpdateForm(BaseModel):
    embedding_model: str


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

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

    try:
Self Denial's avatar
Self Denial committed
149
150
151
        app.state.RAG_EMBEDDING_MODEL_PATH = embedding_model_get_path(
            app.state.RAG_EMBEDDING_MODEL, True
        )
152
153
154
155
156
        app.state.sentence_transformer_ef = (
            embedding_functions.SentenceTransformerEmbeddingFunction(
                model_name=app.state.RAG_EMBEDDING_MODEL_PATH,
                device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
            )
157
        )
Self Denial's avatar
Self Denial committed
158
    except Exception as e:
159
160
161
162
163
164
165
        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
166
        status = False
167

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

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


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


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


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

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


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


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


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


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


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

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


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


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


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

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

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


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

323
324
325
326
327
328
    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
329
330
331
332
333

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


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

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

        return False


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

    return loader, known_type


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

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

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
477
        unvalidated_file_path = f"{UPLOAD_DIR}/{unsanitized_filename}"
478
479
480
481
482
483
484
485
486
487
488
489
        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
490
        filename = file.filename
491
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
492
        with open(file_path, "wb") as f:
493
494
495
            f.write(contents)
            f.close()

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

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

        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
515
516
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
517
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
518
            )
519
    except Exception as e:
520
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
521
522
523
524
525
526
527
528
529
530
        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),
            )
531
532


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

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


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

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

619
        except Exception as e:
620
            log.exception(e)
621
622
623
624

    return True


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


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

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

    return True