main.py 29.3 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
4
5
6
7
8
9
from fastapi import (
    FastAPI,
    Depends,
    HTTPException,
    status,
    UploadFile,
    File,
    Form,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
10
from fastapi.middleware.cors import CORSMiddleware
11
import os, shutil, logging, re
12
13

from pathlib import Path
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
    YOUTUBE_LOADER_LANGUAGE,
97
    RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
98
)
99

100
101
from constants import ERROR_MESSAGES

102
103
104
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
app = FastAPI()

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
115
116
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
117

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

123

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

127
app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
128

Steven Kreitzer's avatar
Steven Kreitzer committed
129

130
app.state.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
131
132
133
app.state.YOUTUBE_LOADER_TRANSLATION = None


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
165
166
167
168
169
170
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
171

Timothy J. Baek's avatar
Timothy J. Baek committed
172
173
174
175
176
177
178
179
180

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
181
182
origins = ["*"]

183

Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
186
187
188
189
190
191
192
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
193
class CollectionNameForm(BaseModel):
194
195
196
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
197
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
198
199
    url: str

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

201
202
203
204
class SearchForm(CollectionNameForm):
    query: str


Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
207
208
209
210
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
211
        "template": app.state.RAG_TEMPLATE,
212
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
213
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
214
        "reranking_model": app.state.RAG_RERANKING_MODEL,
215
216
217
    }


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


Steven Kreitzer's avatar
Steven Kreitzer committed
231
232
233
234
235
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
    return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}


236
237
238
239
240
class OpenAIConfigForm(BaseModel):
    url: str
    key: str


241
class EmbeddingModelUpdateForm(BaseModel):
242
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
243
    embedding_engine: str
244
245
246
    embedding_model: str


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

258
259
        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
260
261
                app.state.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.OPENAI_API_KEY = form_data.openai_config.key
262

263
        update_embedding_model(app.state.RAG_EMBEDDING_MODEL, True)
264

Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
267
268
269
270
271
272
        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,
        )

273
274
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
275
            "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
276
            "embedding_model": app.state.RAG_EMBEDDING_MODEL,
277
            "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
278
279
                "url": app.state.OPENAI_API_BASE_URL,
                "key": app.state.OPENAI_API_KEY,
280
            },
281
282
283
284
285
286
287
        }
    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
288
289


Steven Kreitzer's avatar
Steven Kreitzer committed
290
291
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
292

Steven Kreitzer's avatar
Steven Kreitzer committed
293
294
295
296
297
298
299
300
301
302

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

304
        update_reranking_model(app.state.RAG_RERANKING_MODEL, True)
Steven Kreitzer's avatar
Steven Kreitzer committed
305
306
307
308
309
310
311
312
313
314
315
316
317

        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
318
319
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
320
321
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
322
323
324
325
326
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
327
        "web_loader_ssl_verification": app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
328
329
330
331
        "youtube": {
            "language": app.state.YOUTUBE_LOADER_LANGUAGE,
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
332
333
334
335
336
337
338
339
    }


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


340
341
342
343
344
class YoutubeLoaderConfig(BaseModel):
    language: List[str]
    translation: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
345
class ConfigUpdateForm(BaseModel):
346
347
348
    pdf_extract_images: Optional[bool] = None
    chunk: Optional[ChunkParamUpdateForm] = None
    web_loader_ssl_verification: Optional[bool] = None
349
    youtube: Optional[YoutubeLoaderConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
350
351
352
353


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
    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
375

376
377
378
379
380
381
382
383
384
385
386
387
    app.state.YOUTUBE_LOADER_LANGUAGE = (
        form_data.youtube.language
        if form_data.youtube != None
        else app.state.YOUTUBE_LOADER_LANGUAGE
    )

    app.state.YOUTUBE_LOADER_TRANSLATION = (
        form_data.youtube.translation
        if form_data.youtube != None
        else app.state.YOUTUBE_LOADER_TRANSLATION
    )

Timothy J. Baek's avatar
Timothy J. Baek committed
388
389
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
390
391
392
393
394
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
395
        "web_loader_ssl_verification": app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
396
397
398
399
        "youtube": {
            "language": app.state.YOUTUBE_LOADER_LANGUAGE,
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
400
    }
401
402


Timothy J. Baek's avatar
Timothy J. Baek committed
403
404
405
406
407
408
409
410
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


411
412
413
414
415
416
@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,
417
        "r": app.state.RELEVANCE_THRESHOLD,
Timothy J. Baek's avatar
Timothy J. Baek committed
418
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
419
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
420
421


422
423
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
424
    r: Optional[float] = None
425
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
426
    hybrid: Optional[bool] = None
427
428
429
430
431
432
433
434


@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
435
    app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
Timothy J. Baek's avatar
Timothy J. Baek committed
436
    app.state.ENABLE_RAG_HYBRID_SEARCH = form_data.hybrid if form_data.hybrid else False
Steven Kreitzer's avatar
Steven Kreitzer committed
437
438
439
440
441
    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
442
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
Steven Kreitzer's avatar
Steven Kreitzer committed
443
    }
444
445


446
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
447
448
    collection_name: str
    query: str
449
    k: Optional[int] = None
450
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
451
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
452
453


454
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
455
def query_doc_handler(
456
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
457
458
    user=Depends(get_current_user),
):
459
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
460
461
462
463
        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
464
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
465
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
466
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
467
468
469
470
471
472
                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
473
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
474
475
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
476
    except Exception as e:
477
        log.exception(e)
478
479
480
481
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
482
483


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
484
485
486
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
487
    k: Optional[int] = None
488
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
489
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
490
491


492
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
493
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
494
495
496
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
497
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
498
499
500
501
        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
502
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
503
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
504
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
505
506
507
508
509
510
                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
511
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
512
513
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
514

515
516
517
518
519
520
    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
521
522


Timothy J. Baek's avatar
Timothy J. Baek committed
523
524
525
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
526
527
528
529
530
531
        loader = YoutubeLoader.from_youtube_url(
            form_data.url,
            add_video_info=True,
            language=app.state.YOUTUBE_LOADER_LANGUAGE,
            translation=app.state.YOUTUBE_LOADER_TRANSLATION,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
        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),
        )


552
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
553
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
554
555
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
556
557
558
        loader = get_web_loader(
            form_data.url, verify_ssl=app.state.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
        )
559
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
560
561
562
563
564

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

565
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
566
567
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
568
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
569
570
            "filename": form_data.url,
        }
571
    except Exception as e:
572
        log.exception(e)
573
574
575
576
577
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )

578

579
def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
580
    # Check if the URL is valid
581
    if not validate_url(url):
582
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
583
584
585
586
    return WebBaseLoader(
        url,
        verify_ssl=verify_ssl,
        requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
587
        continue_on_failure=True,
588
    )
589
590


591
592
593
594
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)
595
        if not ENABLE_RAG_LOCAL_WEB_FETCH:
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
            # 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


615
616
617
618
619
620
621
622
623
624
625
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


626
627
628
@app.post("/websearch")
def store_websearch(form_data: SearchForm, user=Depends(get_current_user)):
    try:
629
630
631
632
633
634
635
636
        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,
            )
637
638
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
639
        data = loader.aload()
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658

        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),
        )


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

661
662
663
664
665
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
666

667
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
668
669

    if len(docs) > 0:
670
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
671
672
673
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
674
675
676


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
677
    text, metadata, collection_name, overwrite: bool = False
678
679
680
681
682
683
) -> 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
684
    docs = text_splitter.create_documents([text], metadatas=[metadata])
685
686
687
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

691
692
693
694
695
696
697
    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:
698
                    log.info(f"deleting existing collection {collection_name}")
699
700
                    CHROMA_CLIENT.delete_collection(name=collection_name)

701
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
702

Timothy J. Baek's avatar
Timothy J. Baek committed
703
        embedding_func = get_embedding_function(
Steven Kreitzer's avatar
Steven Kreitzer committed
704
705
706
707
708
709
710
711
            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))
712
        embeddings = embedding_func(embedding_texts)
713
714
715

        for batch in create_batches(
            api=CHROMA_CLIENT,
716
            ids=[str(uuid.uuid4()) for _ in texts],
717
718
719
720
721
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
722

723
        return True
724
    except Exception as e:
725
        log.exception(e)
726
727
728
729
730
731
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


732
733
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
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
    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
782
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
783
784
785
786
787
788
    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)
789
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
790
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
791
792
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
793
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
794
795
        loader = UnstructuredEPubLoader(file_path)
    elif (
796
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
797
798
799
800
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
801
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
802
803
804
805
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
806
807
808
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
809
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
810
    else:
811
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
812
813
814
815
816
        known_type = False

    return loader, known_type


817
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
818
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
819
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
820
821
822
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
823
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
824

825
    log.info(f"file.content_type: {file.content_type}")
826
    try:
827
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
828
        filename = os.path.basename(unsanitized_filename)
829

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

832
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
833
        with open(file_path, "wb") as f:
834
835
836
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
837
838
839
840
841
        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
842
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
843
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
844
845
846
847
848
849
850
851
852
853
854
855

        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
856
857
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
858
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
859
            )
860
    except Exception as e:
861
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
862
863
864
865
866
867
868
869
870
871
        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),
            )
872
873


874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
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
890
891
892
893
894
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
895
896
897
898
899
900
901
902
903
904

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


905
906
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
907
908
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
909
910
911
912
913
914
915
916
917
            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
918
919
920
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
921
922
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
                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,
                                                        )
947
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
948
949
950
951
952
953
954
955
956
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
957
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
958
                    pass
959

960
        except Exception as e:
961
            log.exception(e)
962
963
964
965

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
966
@app.get("/reset/db")
967
968
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
969
970
971


@app.get("/reset")
972
973
974
975
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
976
        try:
977
978
979
980
            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
981
        except Exception as e:
982
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
983

984
985
986
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
987
        log.exception(e)
988
989

    return True