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

from pathlib import Path
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
14
from typing import List
Timothy J. Baek's avatar
Timothy J. Baek committed
15

16
from chromadb.utils.batch_utils import create_batches
Timothy J. Baek's avatar
Timothy J. Baek committed
17

Timothy J. Baek's avatar
Timothy J. Baek committed
18
19
20
21
22
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
23
    BSHTMLLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
24
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
25
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
26
27
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
28
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
29
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
30
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
31
    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
)
Timothy J. Baek's avatar
Timothy J. Baek committed
62

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

71
from config import (
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
72
    ENV,
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
    AppConfig,
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
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
108

109
110
111
112
113
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 = (
114
115
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
Steven Kreitzer's avatar
Steven Kreitzer committed
116

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

120
121
122
123
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
124

125

126
127
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
128

129
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
130

Steven Kreitzer's avatar
Steven Kreitzer committed
131

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


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

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

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

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

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

185

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


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


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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
203
204
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
    return {
        "status": True,
207
208
209
210
211
212
        "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,
213
214
215
    }


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


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


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


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


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

259
        if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
260
            if form_data.openai_config != None:
261
262
                app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.config.OPENAI_API_KEY = form_data.openai_config.key
263

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
264
        update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
265

Timothy J. Baek's avatar
Timothy J. Baek committed
266
        app.state.EMBEDDING_FUNCTION = get_embedding_function(
267
268
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
269
            app.state.sentence_transformer_ef,
270
271
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
272
273
        )

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


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

Steven Kreitzer's avatar
Steven Kreitzer committed
294
295
296
297
298
299

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

305
        update_reranking_model(app.state.config.RAG_RERANKING_MODEL), True
Steven Kreitzer's avatar
Steven Kreitzer committed
306
307
308

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


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


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


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


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
355
356
357
358
    app.state.config.PDF_EXTRACT_IMAGES = (
        form_data.pdf_extract_images
        if form_data.pdf_extract_images is not None
        else app.state.config.PDF_EXTRACT_IMAGES
359
360
    )

361
362
363
364
    app.state.config.CHUNK_SIZE = (
        form_data.chunk.chunk_size
        if form_data.chunk is not None
        else app.state.config.CHUNK_SIZE
365
366
    )

367
368
369
370
    app.state.config.CHUNK_OVERLAP = (
        form_data.chunk.chunk_overlap
        if form_data.chunk is not None
        else app.state.config.CHUNK_OVERLAP
371
372
    )

373
374
375
376
    app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
        form_data.web_loader_ssl_verification
        if form_data.web_loader_ssl_verification != None
        else app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
377
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
378

379
380
381
382
    app.state.config.YOUTUBE_LOADER_LANGUAGE = (
        form_data.youtube.language
        if form_data.youtube is not None
        else app.state.config.YOUTUBE_LOADER_LANGUAGE
383
384
385
386
    )

    app.state.YOUTUBE_LOADER_TRANSLATION = (
        form_data.youtube.translation
387
        if form_data.youtube is not None
388
389
390
        else app.state.YOUTUBE_LOADER_TRANSLATION
    )

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


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


414
415
416
417
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
418
419
420
421
        "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,
422
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
423
424


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


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
436
    app.state.config.RAG_TEMPLATE = (
Timothy J. Baek's avatar
Timothy J. Baek committed
437
        form_data.template if form_data.template else RAG_TEMPLATE
438
    )
439
440
441
    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 = (
Timothy J. Baek's avatar
Timothy J. Baek committed
442
        form_data.hybrid if form_data.hybrid else False
443
    )
Steven Kreitzer's avatar
Steven Kreitzer committed
444
445
    return {
        "status": True,
446
447
448
449
        "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
450
    }
451
452


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


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


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
493
494
495
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
496
    k: Optional[int] = None
497
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
498
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
499
500


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

526
527
528
529
530
531
    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
532
533


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


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

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

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


591
def get_web_loader(url: str, verify_ssl: bool = True):
592
593
594
    # Check if the URL is valid
    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
        # 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)
608
    return WebBaseLoader(url, verify_ssl=verify_ssl)
609
610
611
612
613
614
615
616
617
618
619
620
621


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


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

624
    text_splitter = RecursiveCharacterTextSplitter(
625
626
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
627
628
        add_start_index=True,
    )
629

630
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
631
632

    if len(docs) > 0:
633
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
634
635
636
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
637
638
639


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
640
    text, metadata, collection_name, overwrite: bool = False
641
642
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
643
644
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
645
646
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
647
    docs = text_splitter.create_documents([text], metadatas=[metadata])
648
649
650
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

654
655
656
657
658
659
660
    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:
661
                    log.info(f"deleting existing collection {collection_name}")
662
663
                    CHROMA_CLIENT.delete_collection(name=collection_name)

664
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
665

Timothy J. Baek's avatar
Timothy J. Baek committed
666
        embedding_func = get_embedding_function(
667
668
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
669
            app.state.sentence_transformer_ef,
670
671
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
Steven Kreitzer's avatar
Steven Kreitzer committed
672
673
674
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
675
        embeddings = embedding_func(embedding_texts)
676
677
678

        for batch in create_batches(
            api=CHROMA_CLIENT,
679
            ids=[str(uuid.uuid4()) for _ in texts],
680
681
682
683
684
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
685

686
        return True
687
    except Exception as e:
688
        log.exception(e)
689
690
691
692
693
694
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


695
696
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
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
739
740
741
742
743
744
    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":
745
        loader = PyPDFLoader(
746
            file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES
747
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
748
749
750
751
752
753
    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)
754
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
755
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
756
757
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
758
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
759
760
        loader = UnstructuredEPubLoader(file_path)
    elif (
761
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
762
763
764
765
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
766
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
767
768
769
770
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
771
772
773
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
774
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
775
    else:
776
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
777
778
779
780
781
        known_type = False

    return loader, known_type


782
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
783
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
784
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
785
786
787
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
788
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
789

790
    log.info(f"file.content_type: {file.content_type}")
791
    try:
792
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
793
        filename = os.path.basename(unsanitized_filename)
794

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

797
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
798
        with open(file_path, "wb") as f:
799
800
801
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
802
803
804
805
806
        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
807
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
808
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
809
810
811
812
813
814
815
816
817
818
819
820

        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
821
822
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
823
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
824
            )
825
    except Exception as e:
826
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
827
828
829
830
831
832
833
834
835
836
        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),
            )
837
838


839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
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
855
856
857
858
859
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
860
861
862
863
864
865
866
867
868
869

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


870
871
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
872
873
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
874
875
876
877
878
879
880
881
882
            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
883
884
885
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
886
887
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
                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,
                                                        )
912
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
913
914
915
916
917
918
919
920
921
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
922
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
923
                    pass
924

925
        except Exception as e:
926
            log.exception(e)
927
928
929
930

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
931
@app.get("/reset/db")
932
933
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
934
935
936


@app.get("/reset")
937
938
939
940
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
941
        try:
942
943
944
945
            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
946
        except Exception as e:
947
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
948

949
950
951
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
952
        log.exception(e)
953
954

    return True
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
955
956
957
958
959
960
961
962
963
964
965


if ENV == "dev":

    @app.get("/ef")
    async def get_embeddings():
        return {"result": app.state.EMBEDDING_FUNCTION("hello world")}

    @app.get("/ef/{text}")
    async def get_embeddings_text(text: str):
        return {"result": app.state.EMBEDDING_FUNCTION(text)}