main.py 16.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 sentence_transformers import SentenceTransformer
17
from chromadb.utils import embedding_functions
18
from chromadb.utils.batch_utils import create_batches
Timothy J. Baek's avatar
Timothy J. Baek committed
19

Timothy J. Baek's avatar
Timothy J. Baek committed
20
21
22
23
24
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
25
    BSHTMLLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
26
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
27
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
30
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
31
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
32
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
33
)
34
35
36
37
from langchain.text_splitter import RecursiveCharacterTextSplitter

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

42

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

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

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

70
71
from constants import ERROR_MESSAGES

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
77
app.state.PDF_EXTRACT_IMAGES = False
Timothy J. Baek's avatar
Timothy J. Baek committed
78
79
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
80
app.state.RAG_TEMPLATE = RAG_TEMPLATE
81
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
82
83
app.state.TOP_K = 4

84
85
86
app.state.sentence_transformer_ef = (
    embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name=app.state.RAG_EMBEDDING_MODEL,
87
        device=DEVICE_TYPE,
88
89
    )
)
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
        "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,
142
            device=DEVICE_TYPE,
143
144
145
146
147
        )
    )
    return {
        "status": True,
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
148
149
150
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
151
152
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
153
154
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
155
156
157
158
159
        "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
160
161
162
163
164
165
166
167
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
170
171
172
173
174
175
176
177
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
178
179
180

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
181
182
183
184
185
        "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
186
    }
187
188


Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
191
192
193
194
195
196
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


197
198
199
200
201
202
203
@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
204
205


206
207
208
209
210
211
212
213
214
215
216
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
217
    return {"status": True, "template": app.state.RAG_TEMPLATE}
218
219


220
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
221
222
    collection_name: str
    query: str
223
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
224
225


226
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
227
def query_doc_handler(
228
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
229
230
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
231

232
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
233
234
235
236
        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,
237
238
            embedding_function=app.state.sentence_transformer_ef,
        )
239
    except Exception as e:
240
        log.exception(e)
241
242
243
244
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
245
246


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
247
248
249
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
250
    k: Optional[int] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
251
252


253
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
254
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
255
256
257
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
Timothy J. Baek's avatar
Timothy J. Baek committed
258
259
260
261
262
    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,
263
    )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
264
265


266
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
267
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
268
269
270
271
    # "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
272
273
274
275
276

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

277
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
278
279
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
280
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
281
282
            "filename": form_data.url,
        }
283
    except Exception as e:
284
        log.exception(e)
285
286
287
288
289
290
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


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

293
294
295
296
297
298
    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
299
300
301
302
303

    if len(docs) > 0:
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
304
305
306


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
307
    text, metadata, collection_name, overwrite: bool = False
308
309
310
311
312
313
) -> 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
314
    docs = text_splitter.create_documents([text], metadatas=[metadata])
315
316
317
318
    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
319

320
321
322
323
324
325
326
    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:
327
                    log.info(f"deleting existing collection {collection_name}")
328
329
330
331
332
333
334
                    CHROMA_CLIENT.delete_collection(name=collection_name)

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

335
336
337
338
339
340
341
342
        for batch in create_batches(
            api=CHROMA_CLIENT,
            ids=[str(uuid.uuid1()) for _ in texts],
            metadatas=metadatas,
            documents=texts,
        ):
            collection.add(*batch)

343
344
        return True
    except Exception as e:
345
        log.exception(e)
346
347
348
349
350
351
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


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

    return loader, known_type


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

445
    log.info(f"file.content_type: {file.content_type}")
446
    try:
447
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
448
        filename = os.path.basename(unsanitized_filename)
449

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

452
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
453
        with open(file_path, "wb") as f:
454
455
456
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
457
458
459
460
461
        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
462
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
463
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
464
465
466
467
468
469
470
471
472
473
474
475

        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
476
477
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
478
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
479
            )
480
    except Exception as e:
481
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
482
483
484
485
486
487
488
489
490
491
        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),
            )
492
493


494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
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
510
511
512
513
514
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
515
516
517
518
519
520
521
522
523
524

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


525
526
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
527
528
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
529
530
531
532
533
534
535
536
537
            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
538
539
540
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
541
542
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
                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,
                                                        )
567
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
568
569
570
571
572
573
574
575
576
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
577
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
578
                    pass
579

580
        except Exception as e:
581
            log.exception(e)
582
583
584
585

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
586
@app.get("/reset/db")
587
588
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
589
590
591


@app.get("/reset")
592
593
594
595
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
596
        try:
597
598
599
600
            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
601
        except Exception as e:
602
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
603

604
605
606
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
607
        log.exception(e)
608
609

    return True