main.py 30.2 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
    AppConfig,
99
)
100

101
102
from constants import ERROR_MESSAGES

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

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

108
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
109

110
111
112
113
114
app.state.config.TOP_K = RAG_TOP_K
app.state.config.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD

app.state.config.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
115
116
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
Steven Kreitzer's avatar
Steven Kreitzer committed
117

118
119
app.state.config.CHUNK_SIZE = CHUNK_SIZE
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
120

121
122
123
124
app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
125

126

127
128
app.state.config.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
app.state.config.OPENAI_API_KEY = RAG_OPENAI_API_KEY
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
129

130
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
131

Steven Kreitzer's avatar
Steven Kreitzer committed
132

133
app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
134
135
136
app.state.YOUTUBE_LOADER_TRANSLATION = None


137
138
139
140
def update_embedding_model(
    embedding_model: str,
    update_model: bool = False,
):
141
    if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
        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(
166
    app.state.config.RAG_EMBEDDING_MODEL,
167
168
169
170
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)

update_reranking_model(
171
    app.state.config.RAG_RERANKING_MODEL,
172
173
    RAG_RERANKING_MODEL_AUTO_UPDATE,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
174

Timothy J. Baek's avatar
Timothy J. Baek committed
175
176

app.state.EMBEDDING_FUNCTION = get_embedding_function(
177
178
    app.state.config.RAG_EMBEDDING_ENGINE,
    app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
179
    app.state.sentence_transformer_ef,
180
181
    app.state.config.OPENAI_API_KEY,
    app.state.config.OPENAI_API_BASE_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
182
183
)

Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
origins = ["*"]

186

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


Timothy J. Baek's avatar
Timothy J. Baek committed
196
class CollectionNameForm(BaseModel):
197
198
199
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
200
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
201
202
    url: str

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

204
205
206
207
class SearchForm(CollectionNameForm):
    query: str


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


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


Steven Kreitzer's avatar
Steven Kreitzer committed
234
235
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
236
237
    return {
        "status": True,
238
        "reranking_model": app.state.config.RAG_RERANKING_MODEL,
239
    }
Steven Kreitzer's avatar
Steven Kreitzer committed
240
241


242
243
244
245
246
class OpenAIConfigForm(BaseModel):
    url: str
    key: str


247
class EmbeddingModelUpdateForm(BaseModel):
248
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
249
    embedding_engine: str
250
251
252
    embedding_model: str


Timothy J. Baek's avatar
Timothy J. Baek committed
253
254
@app.post("/embedding/update")
async def update_embedding_config(
255
256
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
Self Denial's avatar
Self Denial committed
257
    log.info(
258
        f"Updating embedding model: {app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
259
    )
260
    try:
261
262
        app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
        app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
Timothy J. Baek's avatar
Timothy J. Baek committed
263

264
        if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
265
            if form_data.openai_config != None:
266
267
                app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.config.OPENAI_API_KEY = form_data.openai_config.key
268

269
        update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL), True
270

Timothy J. Baek's avatar
Timothy J. Baek committed
271
        app.state.EMBEDDING_FUNCTION = get_embedding_function(
272
273
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
274
            app.state.sentence_transformer_ef,
275
276
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
277
278
        )

279
280
        return {
            "status": True,
281
282
            "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
            "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
283
            "openai_config": {
284
285
                "url": app.state.config.OPENAI_API_BASE_URL,
                "key": app.state.config.OPENAI_API_KEY,
286
            },
287
288
289
290
291
292
293
        }
    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
294
295


Steven Kreitzer's avatar
Steven Kreitzer committed
296
297
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
298

Steven Kreitzer's avatar
Steven Kreitzer committed
299
300
301
302
303
304

@app.post("/reranking/update")
async def update_reranking_config(
    form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
):
    log.info(
305
        f"Updating reranking model: {app.state.config.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
Steven Kreitzer's avatar
Steven Kreitzer committed
306
307
    )
    try:
308
        app.state.config.RAG_RERANKING_MODEL = form_data.reranking_model
309

310
        update_reranking_model(app.state.config.RAG_RERANKING_MODEL), True
Steven Kreitzer's avatar
Steven Kreitzer committed
311
312
313

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


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


346
347
348
349
350
class YoutubeLoaderConfig(BaseModel):
    language: List[str]
    translation: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
351
class ConfigUpdateForm(BaseModel):
352
353
354
    pdf_extract_images: Optional[bool] = None
    chunk: Optional[ChunkParamUpdateForm] = None
    web_loader_ssl_verification: Optional[bool] = None
355
    youtube: Optional[YoutubeLoaderConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
356
357
358
359


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
360
    app.state.config.PDF_EXTRACT_IMAGES = (
361
        form_data.pdf_extract_images
362
363
        if form_data.pdf_extract_images is not None
        else app.state.config.PDF_EXTRACT_IMAGES
364
365
    )

366
367
368
369
    app.state.config.CHUNK_SIZE = (
        form_data.chunk.chunk_size
        if form_data.chunk is not None
        else app.state.config.CHUNK_SIZE
370
371
    )

372
    app.state.config.CHUNK_OVERLAP = (
373
        form_data.chunk.chunk_overlap
374
375
        if form_data.chunk is not None
        else app.state.config.CHUNK_OVERLAP
376
377
    )

378
    app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
379
380
        form_data.web_loader_ssl_verification
        if form_data.web_loader_ssl_verification != None
381
        else app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
382
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
383

384
    app.state.config.YOUTUBE_LOADER_LANGUAGE = (
385
        form_data.youtube.language
386
387
        if form_data.youtube is not None
        else app.state.config.YOUTUBE_LOADER_LANGUAGE
388
389
390
391
    )

    app.state.YOUTUBE_LOADER_TRANSLATION = (
        form_data.youtube.translation
392
        if form_data.youtube is not None
393
394
395
        else app.state.YOUTUBE_LOADER_TRANSLATION
    )

Timothy J. Baek's avatar
Timothy J. Baek committed
396
397
    return {
        "status": True,
398
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
399
        "chunk": {
400
401
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
402
        },
403
        "web_loader_ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
404
        "youtube": {
405
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
406
407
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
408
    }
409
410


Timothy J. Baek's avatar
Timothy J. Baek committed
411
412
413
414
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
415
        "template": app.state.config.RAG_TEMPLATE,
Timothy J. Baek's avatar
Timothy J. Baek committed
416
417
418
    }


419
420
421
422
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
423
424
425
426
        "template": app.state.config.RAG_TEMPLATE,
        "k": app.state.config.TOP_K,
        "r": app.state.config.RELEVANCE_THRESHOLD,
        "hybrid": app.state.config.ENABLE_RAG_HYBRID_SEARCH,
427
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
428
429


430
431
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
432
    r: Optional[float] = None
433
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
434
    hybrid: Optional[bool] = None
435
436
437
438
439
440


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
441
    app.state.config.RAG_TEMPLATE = (
442
443
        form_data.template if form_data.template else RAG_TEMPLATE,
    )
444
445
446
    app.state.config.TOP_K = form_data.k if form_data.k else 4
    app.state.config.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
    app.state.config.ENABLE_RAG_HYBRID_SEARCH = (
447
448
        form_data.hybrid if form_data.hybrid else False,
    )
Steven Kreitzer's avatar
Steven Kreitzer committed
449
450
    return {
        "status": True,
451
452
453
454
        "template": app.state.config.RAG_TEMPLATE,
        "k": app.state.config.TOP_K,
        "r": app.state.config.RELEVANCE_THRESHOLD,
        "hybrid": app.state.config.ENABLE_RAG_HYBRID_SEARCH,
Steven Kreitzer's avatar
Steven Kreitzer committed
455
    }
456
457


458
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
459
460
    collection_name: str
    query: str
461
    k: Optional[int] = None
462
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
463
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
464
465


466
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
467
def query_doc_handler(
468
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
469
470
    user=Depends(get_current_user),
):
471
    try:
472
        if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
Timothy J. Baek's avatar
Timothy J. Baek committed
473
474
475
            return query_doc_with_hybrid_search(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
476
                embedding_function=app.state.EMBEDDING_FUNCTION,
477
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
478
                reranking_function=app.state.sentence_transformer_rf,
479
                r=(
480
                    form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
481
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
482
483
484
485
486
            )
        else:
            return query_doc(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
487
                embedding_function=app.state.EMBEDDING_FUNCTION,
488
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
489
            )
490
    except Exception as e:
491
        log.exception(e)
492
493
494
495
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
496
497


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
498
499
500
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
501
    k: Optional[int] = None
502
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
503
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
504
505


506
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
507
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
508
509
510
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
511
    try:
512
        if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
Timothy J. Baek's avatar
Timothy J. Baek committed
513
514
515
            return query_collection_with_hybrid_search(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
516
                embedding_function=app.state.EMBEDDING_FUNCTION,
517
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
518
                reranking_function=app.state.sentence_transformer_rf,
519
                r=(
520
                    form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
521
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
522
523
524
525
526
            )
        else:
            return query_collection(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
527
                embedding_function=app.state.EMBEDDING_FUNCTION,
528
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
529
            )
530

531
532
533
534
535
536
    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
537
538


Timothy J. Baek's avatar
Timothy J. Baek committed
539
540
541
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
542
543
544
        loader = YoutubeLoader.from_youtube_url(
            form_data.url,
            add_video_info=True,
545
            language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
546
547
            translation=app.state.YOUTUBE_LOADER_TRANSLATION,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
        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),
        )


568
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
569
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
570
571
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
572
        loader = get_web_loader(
573
            form_data.url,
574
            verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
575
        )
576
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
577
578
579
580
581

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

582
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
583
584
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
585
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
586
587
            "filename": form_data.url,
        }
588
    except Exception as e:
589
        log.exception(e)
590
591
592
593
594
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )

595

596
def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
597
    # Check if the URL is valid
598
    if not validate_url(url):
599
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
600
601
602
603
    return WebBaseLoader(
        url,
        verify_ssl=verify_ssl,
        requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
604
        continue_on_failure=True,
605
    )
606
607


608
609
610
611
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)
612
        if not ENABLE_RAG_LOCAL_WEB_FETCH:
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
            # 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


632
633
634
635
636
637
638
639
640
641
642
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


643
644
645
@app.post("/websearch")
def store_websearch(form_data: SearchForm, user=Depends(get_current_user)):
    try:
646
647
648
649
650
651
652
653
        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,
            )
654
655
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
656
        data = loader.aload()
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675

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


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

678
    text_splitter = RecursiveCharacterTextSplitter(
679
680
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
681
682
        add_start_index=True,
    )
683

684
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
685
686

    if len(docs) > 0:
687
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
688
689
690
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
691
692
693


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
694
    text, metadata, collection_name, overwrite: bool = False
695
696
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
697
698
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
699
700
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
701
    docs = text_splitter.create_documents([text], metadatas=[metadata])
702
703
704
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

708
709
710
711
712
713
714
    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:
715
                    log.info(f"deleting existing collection {collection_name}")
716
717
                    CHROMA_CLIENT.delete_collection(name=collection_name)

718
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
719

Timothy J. Baek's avatar
Timothy J. Baek committed
720
        embedding_func = get_embedding_function(
721
722
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
723
            app.state.sentence_transformer_ef,
724
725
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
Steven Kreitzer's avatar
Steven Kreitzer committed
726
727
728
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
729
        embeddings = embedding_func(embedding_texts)
730
731
732

        for batch in create_batches(
            api=CHROMA_CLIENT,
733
            ids=[str(uuid.uuid4()) for _ in texts],
734
735
736
737
738
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
739

740
        return True
741
    except Exception as e:
742
        log.exception(e)
743
744
745
746
747
748
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


749
750
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
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
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
    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":
799
        loader = PyPDFLoader(
800
            file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES
801
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
802
803
804
805
806
807
    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)
808
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
809
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
810
811
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
812
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
813
814
        loader = UnstructuredEPubLoader(file_path)
    elif (
815
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
816
817
818
819
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
820
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
821
822
823
824
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
825
826
827
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
828
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
829
    else:
830
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
831
832
833
834
835
        known_type = False

    return loader, known_type


836
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
837
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
838
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
839
840
841
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
842
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
843

844
    log.info(f"file.content_type: {file.content_type}")
845
    try:
846
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
847
        filename = os.path.basename(unsanitized_filename)
848

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

851
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
852
        with open(file_path, "wb") as f:
853
854
855
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
856
857
858
859
860
        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
861
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
862
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
863
864
865
866
867
868
869
870
871
872
873
874

        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
875
876
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
877
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
878
            )
879
    except Exception as e:
880
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
881
882
883
884
885
886
887
888
889
890
        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),
            )
891
892


893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
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
909
910
911
912
913
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
914
915
916
917
918
919
920
921
922
923

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


924
925
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
926
927
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
928
929
930
931
932
933
934
935
936
            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
937
938
939
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
940
941
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
                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,
                                                        )
966
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
967
968
969
970
971
972
973
974
975
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
976
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
977
                    pass
978

979
        except Exception as e:
980
            log.exception(e)
981
982
983
984

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
985
@app.get("/reset/db")
986
987
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
988
989
990


@app.get("/reset")
991
992
993
994
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
995
        try:
996
997
998
999
            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
1000
        except Exception as e:
1001
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
1002

1003
1004
1005
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
1006
        log.exception(e)
1007
1008

    return True