main.py 22 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.batch_utils import create_batches
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
import sentence_transformers
41

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

48
from apps.rag.utils import (
49
    get_model_path,
50
    query_embeddings_doc,
Steven Kreitzer's avatar
Steven Kreitzer committed
51
    query_embeddings_function,
52
53
    query_embeddings_collection,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
54

55
56
57
58
59
60
from utils.misc import (
    calculate_sha256,
    calculate_sha256_string,
    sanitize_filename,
    extract_folders_after_data_docs,
)
61
from utils.utils import get_current_user, get_admin_user
62

63
from config import (
64
    SRC_LOG_LEVELS,
65
66
    UPLOAD_DIR,
    DOCS_DIR,
67
68
    RAG_TOP_K,
    RAG_RELEVANCE_THRESHOLD,
69
    RAG_EMBEDDING_ENGINE,
70
    RAG_EMBEDDING_MODEL,
71
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
72
    RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
Steven Kreitzer's avatar
Steven Kreitzer committed
73
    RAG_RERANKING_MODEL,
74
    RAG_RERANKING_MODEL_AUTO_UPDATE,
Steven Kreitzer's avatar
Steven Kreitzer committed
75
    RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
76
77
    RAG_OPENAI_API_BASE_URL,
    RAG_OPENAI_API_KEY,
78
    DEVICE_TYPE,
79
80
81
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
82
    RAG_TEMPLATE,
83
)
84

85
86
from constants import ERROR_MESSAGES

87
88
89
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
90
91
app = FastAPI()

92
93
app.state.TOP_K = RAG_TOP_K
app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
Timothy J. Baek's avatar
Timothy J. Baek committed
94
95
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
96

97
app.state.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
98
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
Steven Kreitzer's avatar
Steven Kreitzer committed
99
app.state.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
100
app.state.RAG_TEMPLATE = RAG_TEMPLATE
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
101

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
102
103
app.state.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = RAG_OPENAI_API_KEY
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
104

105
106
app.state.PDF_EXTRACT_IMAGES = False

Steven Kreitzer's avatar
Steven Kreitzer committed
107

108
109
110
111
112
113
114
115
116
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
142
143
144
def update_embedding_model(
    embedding_model: str,
    update_model: bool = False,
):
    if embedding_model and app.state.RAG_EMBEDDING_ENGINE == "":
        app.state.sentence_transformer_ef = sentence_transformers.SentenceTransformer(
            get_model_path(embedding_model, update_model),
            device=DEVICE_TYPE,
            trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
        )
    else:
        app.state.sentence_transformer_ef = None


def update_reranking_model(
    reranking_model: str,
    update_model: bool = False,
):
    if reranking_model:
        app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
            get_model_path(reranking_model, update_model),
            device=DEVICE_TYPE,
            trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
        )
    else:
        app.state.sentence_transformer_rf = None


update_embedding_model(
    app.state.RAG_EMBEDDING_MODEL,
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)

update_reranking_model(
    app.state.RAG_RERANKING_MODEL,
    RAG_RERANKING_MODEL_AUTO_UPDATE,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
145

Timothy J. Baek's avatar
Timothy J. Baek committed
146
147
origins = ["*"]

148

Timothy J. Baek's avatar
Timothy J. Baek committed
149
150
151
152
153
154
155
156
157
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
158
class CollectionNameForm(BaseModel):
159
160
161
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
162
163
164
class StoreWebForm(CollectionNameForm):
    url: str

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

Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
170
171
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
172
        "template": app.state.RAG_TEMPLATE,
173
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
174
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
175
        "reranking_model": app.state.RAG_RERANKING_MODEL,
176
177
178
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
179
180
@app.get("/embedding")
async def get_embedding_config(user=Depends(get_admin_user)):
181
182
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
183
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
184
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
185
        "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
186
187
            "url": app.state.OPENAI_API_BASE_URL,
            "key": app.state.OPENAI_API_KEY,
188
        },
189
190
191
    }


Steven Kreitzer's avatar
Steven Kreitzer committed
192
193
194
195
196
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
    return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}


197
198
199
200
201
class OpenAIConfigForm(BaseModel):
    url: str
    key: str


202
class EmbeddingModelUpdateForm(BaseModel):
203
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
204
    embedding_engine: str
205
206
207
    embedding_model: str


Timothy J. Baek's avatar
Timothy J. Baek committed
208
209
@app.post("/embedding/update")
async def update_embedding_config(
210
211
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
Self Denial's avatar
Self Denial committed
212
213
    log.info(
        f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
214
    )
215
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
216
        app.state.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
217
        app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
Timothy J. Baek's avatar
Timothy J. Baek committed
218

219
220
        if app.state.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
            if form_data.openai_config != None:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
221
222
                app.state.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.OPENAI_API_KEY = form_data.openai_config.key
223

224
        update_embedding_model(app.state.RAG_EMBEDDING_MODEL, True)
225

226
227
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
228
            "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
229
            "embedding_model": app.state.RAG_EMBEDDING_MODEL,
230
            "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
231
232
                "url": app.state.OPENAI_API_BASE_URL,
                "key": app.state.OPENAI_API_KEY,
233
            },
234
235
236
237
238
239
240
        }
    except Exception as e:
        log.exception(f"Problem updating embedding model: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
241
242


Steven Kreitzer's avatar
Steven Kreitzer committed
243
244
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
245

Steven Kreitzer's avatar
Steven Kreitzer committed
246
247
248
249
250
251
252
253
254
255

@app.post("/reranking/update")
async def update_reranking_config(
    form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
):
    log.info(
        f"Updating reranking model: {app.state.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
    )
    try:
        app.state.RAG_RERANKING_MODEL = form_data.reranking_model
256

257
        update_reranking_model(app.state.RAG_RERANKING_MODEL, True)
Steven Kreitzer's avatar
Steven Kreitzer committed
258
259
260
261
262
263
264
265
266
267
268
269
270

        return {
            "status": True,
            "reranking_model": app.state.RAG_RERANKING_MODEL,
        }
    except Exception as e:
        log.exception(f"Problem updating reranking model: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
271
272
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
273
274
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
275
276
277
278
279
        "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
280
281
282
283
284
285
286
287
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
288
289
290
291
292
293
294
295
296
297
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
298
299
300

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
301
302
303
304
305
        "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
306
    }
307
308


Timothy J. Baek's avatar
Timothy J. Baek committed
309
310
311
312
313
314
315
316
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


317
318
319
320
321
322
@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,
323
        "r": app.state.RELEVANCE_THRESHOLD,
324
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
325
326


327
328
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
329
    r: Optional[float] = None
330
331
332
333
334
335
336
337
338
    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
339
    app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
Timothy J. Baek's avatar
Timothy J. Baek committed
340
    return {"status": True, "template": app.state.RAG_TEMPLATE}
341
342


343
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
344
345
    collection_name: str
    query: str
346
    k: Optional[int] = None
347
    r: Optional[float] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
348
349


350
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
351
def query_doc_handler(
352
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
353
354
    user=Depends(get_current_user),
):
355
    try:
Steven Kreitzer's avatar
Steven Kreitzer committed
356
357
358
359
360
361
362
        embeddings_function = query_embeddings_function(
            app.state.RAG_EMBEDDING_ENGINE,
            app.state.RAG_EMBEDDING_MODEL,
            app.state.sentence_transformer_ef,
            app.state.OPENAI_API_KEY,
            app.state.OPENAI_API_BASE_URL,
        )
363

364
365
366
367
        return query_embeddings_doc(
            collection_name=form_data.collection_name,
            query=form_data.query,
            k=form_data.k if form_data.k else app.state.TOP_K,
368
            r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
Steven Kreitzer's avatar
Steven Kreitzer committed
369
370
            embeddings_function=embeddings_function,
            reranking_function=app.state.sentence_transformer_rf,
371
        )
372
    except Exception as e:
373
        log.exception(e)
374
375
376
377
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
378
379


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
380
381
382
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
383
    k: Optional[int] = None
384
    r: Optional[float] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
385
386


387
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
388
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
389
390
391
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
392
    try:
Steven Kreitzer's avatar
Steven Kreitzer committed
393
394
395
396
397
398
399
        embeddings_function = embeddings_function(
            app.state.RAG_EMBEDDING_ENGINE,
            app.state.RAG_EMBEDDING_MODEL,
            app.state.sentence_transformer_ef,
            app.state.OPENAI_API_KEY,
            app.state.OPENAI_API_BASE_URL,
        )
400

401
402
        return query_embeddings_collection(
            collection_names=form_data.collection_names,
Steven Kreitzer's avatar
Steven Kreitzer committed
403
            query=form_data.query,
404
            k=form_data.k if form_data.k else app.state.TOP_K,
405
            r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
Steven Kreitzer's avatar
Steven Kreitzer committed
406
407
            embeddings_function=embeddings_function,
            reranking_function=app.state.sentence_transformer_rf,
408
        )
409
410
411
412
413
414
    except Exception as e:
        log.exception(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
415
416


417
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
418
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
419
420
421
422
    # "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
423
424
425
426
427

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

428
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
429
430
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
431
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
432
433
            "filename": form_data.url,
        }
434
    except Exception as e:
435
        log.exception(e)
436
437
438
439
440
441
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


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

444
445
446
447
448
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
449

450
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
451
452

    if len(docs) > 0:
453
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
454
455
456
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
457
458
459


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
460
    text, metadata, collection_name, overwrite: bool = False
461
462
463
464
465
466
) -> 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
467
    docs = text_splitter.create_documents([text], metadatas=[metadata])
468
469
470
    return store_docs_in_vector_db(docs, collection_name, overwrite)


Timothy J. Baek's avatar
Timothy J. Baek committed
471
def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> bool:
472
    log.info(f"store_docs_in_vector_db {docs} {collection_name}")
Timothy J. Baek's avatar
Timothy J. Baek committed
473

474
475
476
477
478
479
480
    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:
481
                    log.info(f"deleting existing collection {collection_name}")
482
483
                    CHROMA_CLIENT.delete_collection(name=collection_name)

484
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
485

Steven Kreitzer's avatar
Steven Kreitzer committed
486
487
488
489
490
491
492
493
494
        embedding_func = query_embeddings_function(
            app.state.RAG_EMBEDDING_ENGINE,
            app.state.RAG_EMBEDDING_MODEL,
            app.state.sentence_transformer_ef,
            app.state.OPENAI_API_KEY,
            app.state.OPENAI_API_BASE_URL,
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
495
        embeddings = embedding_func(embedding_texts)
496
497
498
499
500
501
502
503
504

        for batch in create_batches(
            api=CHROMA_CLIENT,
            ids=[str(uuid.uuid1()) for _ in texts],
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
505

506
        return True
507
    except Exception as e:
508
        log.exception(e)
509
510
511
512
513
514
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


515
516
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
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
    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
565
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
566
567
568
569
570
571
    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)
572
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
573
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
574
575
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
576
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
577
578
        loader = UnstructuredEPubLoader(file_path)
    elif (
579
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
580
581
582
583
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
584
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
585
586
587
588
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
589
590
591
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
592
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
593
    else:
594
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
595
596
597
598
599
        known_type = False

    return loader, known_type


600
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
601
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
602
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
603
604
605
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
606
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
607

608
    log.info(f"file.content_type: {file.content_type}")
609
    try:
610
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
611
        filename = os.path.basename(unsanitized_filename)
612

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

615
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
616
        with open(file_path, "wb") as f:
617
618
619
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
620
621
622
623
624
        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
625
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
626
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
627
628
629
630
631
632
633
634
635
636
637
638

        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
639
640
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
641
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
642
            )
643
    except Exception as e:
644
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
645
646
647
648
649
650
651
652
653
654
        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),
            )
655
656


657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
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
673
674
675
676
677
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
678
679
680
681
682
683
684
685
686
687

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


688
689
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
690
691
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
692
693
694
695
696
697
698
699
700
            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
701
702
703
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
704
705
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
                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,
                                                        )
730
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
731
732
733
734
735
736
737
738
739
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
740
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
741
                    pass
742

743
        except Exception as e:
744
            log.exception(e)
745
746
747
748

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
749
@app.get("/reset/db")
750
751
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
752
753
754


@app.get("/reset")
755
756
757
758
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
759
        try:
760
761
762
763
            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
764
        except Exception as e:
765
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
766

767
768
769
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
770
        log.exception(e)
771
772

    return True