main.py 28 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
14
from typing import List, Union, Sequence
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
    YoutubeLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
32
)
33
34
from langchain.text_splitter import RecursiveCharacterTextSplitter

35
36
37
38
39
import validators
import urllib.parse
import socket


40
41
from pydantic import BaseModel
from typing import Optional
42
import mimetypes
43
import uuid
44
45
import json

46
import sentence_transformers
47

48
49
50
51
52
from apps.web.models.documents import (
    Documents,
    DocumentForm,
    DocumentResponse,
)
Jannik Streidl's avatar
Jannik Streidl committed
53

54
from apps.rag.utils import (
55
    get_model_path,
Timothy J. Baek's avatar
Timothy J. Baek committed
56
57
58
59
60
    get_embedding_function,
    query_doc,
    query_doc_with_hybrid_search,
    query_collection,
    query_collection_with_hybrid_search,
61
    search_web,
62
)
Timothy J. Baek's avatar
Timothy J. Baek committed
63

64
65
66
67
68
69
from utils.misc import (
    calculate_sha256,
    calculate_sha256_string,
    sanitize_filename,
    extract_folders_after_data_docs,
)
70
from utils.utils import get_current_user, get_admin_user
71

72
from config import (
73
    SRC_LOG_LEVELS,
74
75
    UPLOAD_DIR,
    DOCS_DIR,
76
77
    RAG_TOP_K,
    RAG_RELEVANCE_THRESHOLD,
78
    RAG_EMBEDDING_ENGINE,
79
    RAG_EMBEDDING_MODEL,
80
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
81
    RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
82
    ENABLE_RAG_HYBRID_SEARCH,
83
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Steven Kreitzer's avatar
Steven Kreitzer committed
84
    RAG_RERANKING_MODEL,
85
    PDF_EXTRACT_IMAGES,
86
    RAG_RERANKING_MODEL_AUTO_UPDATE,
Steven Kreitzer's avatar
Steven Kreitzer committed
87
    RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
88
89
    RAG_OPENAI_API_BASE_URL,
    RAG_OPENAI_API_KEY,
90
    DEVICE_TYPE,
91
92
93
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
94
    RAG_TEMPLATE,
95
    ENABLE_RAG_LOCAL_WEB_FETCH,
96
)
97

98
99
from constants import ERROR_MESSAGES

100
101
102
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
103
104
app = FastAPI()

105
106
app.state.TOP_K = RAG_TOP_K
app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
Timothy J. Baek's avatar
Timothy J. Baek committed
107
108

app.state.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
109
110
111
app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
Steven Kreitzer's avatar
Steven Kreitzer committed
112

Timothy J. Baek's avatar
Timothy J. Baek committed
113
114
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
115

116
app.state.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
117
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
Steven Kreitzer's avatar
Steven Kreitzer committed
118
app.state.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
119
app.state.RAG_TEMPLATE = RAG_TEMPLATE
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
120

121

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
122
123
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
124

125
app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
126

Steven Kreitzer's avatar
Steven Kreitzer committed
127

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
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
165

Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
168
169
170
171
172
173
174

app.state.EMBEDDING_FUNCTION = get_embedding_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,
)

Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
origins = ["*"]

177

Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
180
181
182
183
184
185
186
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
187
class CollectionNameForm(BaseModel):
188
189
190
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
191
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
192
193
    url: str

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

195
196
197
198
class SearchForm(CollectionNameForm):
    query: str


Timothy J. Baek's avatar
Timothy J. Baek committed
199
200
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
201
202
203
204
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
205
        "template": app.state.RAG_TEMPLATE,
206
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
207
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
208
        "reranking_model": app.state.RAG_RERANKING_MODEL,
209
210
211
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
212
213
@app.get("/embedding")
async def get_embedding_config(user=Depends(get_admin_user)):
214
215
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
216
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
217
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
218
        "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
219
220
            "url": app.state.OPENAI_API_BASE_URL,
            "key": app.state.OPENAI_API_KEY,
221
        },
222
223
224
    }


Steven Kreitzer's avatar
Steven Kreitzer committed
225
226
227
228
229
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
    return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}


230
231
232
233
234
class OpenAIConfigForm(BaseModel):
    url: str
    key: str


235
class EmbeddingModelUpdateForm(BaseModel):
236
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
237
    embedding_engine: str
238
239
240
    embedding_model: str


Timothy J. Baek's avatar
Timothy J. Baek committed
241
242
@app.post("/embedding/update")
async def update_embedding_config(
243
244
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
Self Denial's avatar
Self Denial committed
245
246
    log.info(
        f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
247
    )
248
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
249
        app.state.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
250
        app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
Timothy J. Baek's avatar
Timothy J. Baek committed
251

252
253
        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
254
255
                app.state.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.OPENAI_API_KEY = form_data.openai_config.key
256

257
        update_embedding_model(app.state.RAG_EMBEDDING_MODEL, True)
258

Timothy J. Baek's avatar
Timothy J. Baek committed
259
260
261
262
263
264
265
266
        app.state.EMBEDDING_FUNCTION = get_embedding_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,
        )

267
268
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
269
            "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
270
            "embedding_model": app.state.RAG_EMBEDDING_MODEL,
271
            "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
272
273
                "url": app.state.OPENAI_API_BASE_URL,
                "key": app.state.OPENAI_API_KEY,
274
            },
275
276
277
278
279
280
281
        }
    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
282
283


Steven Kreitzer's avatar
Steven Kreitzer committed
284
285
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
286

Steven Kreitzer's avatar
Steven Kreitzer committed
287
288
289
290
291
292
293
294
295
296

@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
297

298
        update_reranking_model(app.state.RAG_RERANKING_MODEL, True)
Steven Kreitzer's avatar
Steven Kreitzer committed
299
300
301
302
303
304
305
306
307
308
309
310
311

        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
312
313
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
314
315
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
316
317
318
319
320
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
321
        "web_loader_ssl_verification": app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Timothy J. Baek's avatar
Timothy J. Baek committed
322
323
324
325
326
327
328
329
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
330
class ConfigUpdateForm(BaseModel):
331
332
333
    pdf_extract_images: Optional[bool] = None
    chunk: Optional[ChunkParamUpdateForm] = None
    web_loader_ssl_verification: Optional[bool] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
334
335
336
337


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
    app.state.PDF_EXTRACT_IMAGES = (
        form_data.pdf_extract_images
        if form_data.pdf_extract_images != None
        else app.state.PDF_EXTRACT_IMAGES
    )

    app.state.CHUNK_SIZE = (
        form_data.chunk.chunk_size if form_data.chunk != None else app.state.CHUNK_SIZE
    )

    app.state.CHUNK_OVERLAP = (
        form_data.chunk.chunk_overlap
        if form_data.chunk != None
        else app.state.CHUNK_OVERLAP
    )

    app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
        form_data.web_loader_ssl_verification
        if form_data.web_loader_ssl_verification != None
        else app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
359
360
361

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
362
363
364
365
366
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
367
        "web_loader_ssl_verification": app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Timothy J. Baek's avatar
Timothy J. Baek committed
368
    }
369
370


Timothy J. Baek's avatar
Timothy J. Baek committed
371
372
373
374
375
376
377
378
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


379
380
381
382
383
384
@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,
385
        "r": app.state.RELEVANCE_THRESHOLD,
Timothy J. Baek's avatar
Timothy J. Baek committed
386
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
387
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
388
389


390
391
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
392
    r: Optional[float] = None
393
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
394
    hybrid: Optional[bool] = None
395
396
397
398
399
400
401
402


@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
403
    app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
Timothy J. Baek's avatar
Timothy J. Baek committed
404
    app.state.ENABLE_RAG_HYBRID_SEARCH = form_data.hybrid if form_data.hybrid else False
Steven Kreitzer's avatar
Steven Kreitzer committed
405
406
407
408
409
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
        "k": app.state.TOP_K,
        "r": app.state.RELEVANCE_THRESHOLD,
Timothy J. Baek's avatar
Timothy J. Baek committed
410
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
Steven Kreitzer's avatar
Steven Kreitzer committed
411
    }
412
413


414
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
415
416
    collection_name: str
    query: str
417
    k: Optional[int] = None
418
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
419
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
420
421


422
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
423
def query_doc_handler(
424
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
425
426
    user=Depends(get_current_user),
):
427
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
428
429
430
431
        if app.state.ENABLE_RAG_HYBRID_SEARCH:
            return query_doc_with_hybrid_search(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
432
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
433
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
434
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
435
436
437
438
439
440
                r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
            )
        else:
            return query_doc(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
441
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
442
443
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
444
    except Exception as e:
445
        log.exception(e)
446
447
448
449
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
450
451


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
452
453
454
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
455
    k: Optional[int] = None
456
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
457
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
458
459


460
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
461
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
462
463
464
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
465
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
466
467
468
469
        if app.state.ENABLE_RAG_HYBRID_SEARCH:
            return query_collection_with_hybrid_search(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
470
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
471
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
472
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
473
474
475
476
477
478
                r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
            )
        else:
            return query_collection(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
479
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
480
481
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
482

483
484
485
486
487
488
    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
489
490


Timothy J. Baek's avatar
Timothy J. Baek committed
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
        loader = YoutubeLoader.from_youtube_url(form_data.url, add_video_info=False)
        data = loader.load()

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

        store_data_in_vector_db(data, collection_name, overwrite=True)
        return {
            "status": True,
            "collection_name": collection_name,
            "filename": form_data.url,
        }
    except Exception as e:
        log.exception(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


515
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
516
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
517
518
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
519
520
521
        loader = get_web_loader(
            form_data.url, verify_ssl=app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
        )
522
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
523
524
525
526
527

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

528
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
529
530
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
531
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
532
533
            "filename": form_data.url,
        }
534
    except Exception as e:
535
        log.exception(e)
536
537
538
539
540
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )

541
def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
542
    # Check if the URL is valid
543
    if not validate_url(url):
544
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
545
    return WebBaseLoader(url, verify_ssl=verify_ssl)
546
547


548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
def validate_url(url: Union[str, Sequence[str]]):
    if isinstance(url, str):
        if isinstance(validators.url(url), validators.ValidationError):
            raise ValueError(ERROR_MESSAGES.INVALID_URL)
        if not ENABLE_LOCAL_WEB_FETCH:
            # Local web fetch is disabled, filter out any URLs that resolve to private IP addresses
            parsed_url = urllib.parse.urlparse(url)
            # Get IPv4 and IPv6 addresses
            ipv4_addresses, ipv6_addresses = resolve_hostname(parsed_url.hostname)
            # Check if any of the resolved addresses are private
            # This is technically still vulnerable to DNS rebinding attacks, as we don't control WebBaseLoader
            for ip in ipv4_addresses:
                if validators.ipv4(ip, private=True):
                    raise ValueError(ERROR_MESSAGES.INVALID_URL)
            for ip in ipv6_addresses:
                if validators.ipv6(ip, private=True):
                    raise ValueError(ERROR_MESSAGES.INVALID_URL)
        return True
    elif isinstance(url, Sequence):
        return all(validate_url(u) for u in url)
    else:
        return False


572
573
574
575
576
577
578
579
580
581
582
def resolve_hostname(hostname):
    # Get address information
    addr_info = socket.getaddrinfo(hostname, None)

    # Extract IP addresses from address information
    ipv4_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET]
    ipv6_addresses = [info[4][0] for info in addr_info if info[0] == socket.AF_INET6]

    return ipv4_addresses, ipv6_addresses


583
584
585
@app.post("/websearch")
def store_websearch(form_data: SearchForm, user=Depends(get_current_user)):
    try:
586
587
588
589
590
591
592
593
        try:
            web_results = search_web(form_data.query)
        except Exception as e:
            log.exception(e)
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.WEB_SEARCH_ERROR,
            )
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
        data = loader.load()

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

        store_data_in_vector_db(data, collection_name, overwrite=True)
        return {
            "status": True,
            "collection_name": collection_name,
            "filenames": urls,
        }
    except Exception as e:
        log.exception(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


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

618
619
620
621
622
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
623

624
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
625
626

    if len(docs) > 0:
627
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
628
629
630
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
631
632
633


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
634
    text, metadata, collection_name, overwrite: bool = False
635
636
637
638
639
640
) -> 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
641
    docs = text_splitter.create_documents([text], metadatas=[metadata])
642
643
644
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

648
649
650
651
652
653
654
    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:
655
                    log.info(f"deleting existing collection {collection_name}")
656
657
                    CHROMA_CLIENT.delete_collection(name=collection_name)

658
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
659

Timothy J. Baek's avatar
Timothy J. Baek committed
660
        embedding_func = get_embedding_function(
Steven Kreitzer's avatar
Steven Kreitzer committed
661
662
663
664
665
666
667
668
            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))
669
        embeddings = embedding_func(embedding_texts)
670
671
672
673
674
675
676
677
678

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

680
        return True
681
    except Exception as e:
682
        log.exception(e)
683
684
685
686
687
688
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


689
690
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
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
    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
739
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
740
741
742
743
744
745
    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)
746
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
747
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
748
749
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
750
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
751
752
        loader = UnstructuredEPubLoader(file_path)
    elif (
753
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
754
755
756
757
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
758
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
759
760
761
762
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
763
764
765
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
766
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
767
    else:
768
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
769
770
771
772
773
        known_type = False

    return loader, known_type


774
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
775
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
776
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
777
778
779
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
780
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
781

782
    log.info(f"file.content_type: {file.content_type}")
783
    try:
784
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
785
        filename = os.path.basename(unsanitized_filename)
786

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

789
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
790
        with open(file_path, "wb") as f:
791
792
793
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
794
795
796
797
798
        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
799
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
800
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
801
802
803
804
805
806
807
808
809
810
811
812

        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
813
814
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
815
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
816
            )
817
    except Exception as e:
818
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
819
820
821
822
823
824
825
826
827
828
        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),
            )
829
830


831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
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
847
848
849
850
851
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
852
853
854
855
856
857
858
859
860
861

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


862
863
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
864
865
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
866
867
868
869
870
871
872
873
874
            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
875
876
877
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
878
879
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
                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,
                                                        )
904
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
905
906
907
908
909
910
911
912
913
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
914
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
915
                    pass
916

917
        except Exception as e:
918
            log.exception(e)
919
920
921
922

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
923
@app.get("/reset/db")
924
925
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
926
927
928


@app.get("/reset")
929
930
931
932
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
933
        try:
934
935
936
937
            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
938
        except Exception as e:
939
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
940

941
942
943
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
944
        log.exception(e)
945
946

    return True