main.py 41.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
mindspawn's avatar
mindspawn committed
12
from datetime import datetime
13
14

from pathlib import Path
15
from typing import List, Union, Sequence, Iterator, Any
Timothy J. Baek's avatar
Timothy J. Baek committed
16

17
from chromadb.utils.batch_utils import create_batches
18
from langchain_core.documents import Document
Timothy J. Baek's avatar
Timothy J. Baek committed
19

Timothy J. Baek's avatar
Timothy J. Baek committed
20
21
22
23
24
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
25
    BSHTMLLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
26
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
27
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
30
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
31
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
32
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
33
    UnstructuredPowerPointLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
34
    YoutubeLoader,
mindspawn's avatar
mindspawn committed
35
    OutlookMessageLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
36
)
37
38
from langchain.text_splitter import RecursiveCharacterTextSplitter

39
40
41
42
43
import validators
import urllib.parse
import socket


44
45
from pydantic import BaseModel
from typing import Optional
46
import mimetypes
47
import uuid
48
49
import json

50
import sentence_transformers
51

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
52
from apps.webui.models.documents import (
53
54
55
56
    Documents,
    DocumentForm,
    DocumentResponse,
)
Jannik Streidl's avatar
Jannik Streidl committed
57

58
from apps.rag.utils import (
59
    get_model_path,
Timothy J. Baek's avatar
Timothy J. Baek committed
60
61
62
63
64
    get_embedding_function,
    query_doc,
    query_doc_with_hybrid_search,
    query_collection,
    query_collection_with_hybrid_search,
65
)
Timothy J. Baek's avatar
Timothy J. Baek committed
66

Timothy J. Baek's avatar
Timothy J. Baek committed
67
68
69
70
71
72
from apps.rag.search.brave import search_brave
from apps.rag.search.google_pse import search_google_pse
from apps.rag.search.main import SearchResult
from apps.rag.search.searxng import search_searxng
from apps.rag.search.serper import search_serper
from apps.rag.search.serpstack import search_serpstack
73
from apps.rag.search.serply import search_serply
74
from apps.rag.search.duckduckgo import search_duckduckgo
Timothy J. Baek's avatar
Timothy J. Baek committed
75

76
77
78
79
80
81
from utils.misc import (
    calculate_sha256,
    calculate_sha256_string,
    sanitize_filename,
    extract_folders_after_data_docs,
)
82
from utils.utils import get_current_user, get_admin_user
83

84
from config import (
85
    AppConfig,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
86
    ENV,
87
    SRC_LOG_LEVELS,
88
89
    UPLOAD_DIR,
    DOCS_DIR,
90
91
    RAG_TOP_K,
    RAG_RELEVANCE_THRESHOLD,
92
    RAG_EMBEDDING_ENGINE,
93
    RAG_EMBEDDING_MODEL,
94
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
95
    RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
96
    ENABLE_RAG_HYBRID_SEARCH,
97
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Steven Kreitzer's avatar
Steven Kreitzer committed
98
    RAG_RERANKING_MODEL,
99
    PDF_EXTRACT_IMAGES,
100
    RAG_RERANKING_MODEL_AUTO_UPDATE,
Steven Kreitzer's avatar
Steven Kreitzer committed
101
    RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
102
103
    RAG_OPENAI_API_BASE_URL,
    RAG_OPENAI_API_KEY,
104
    DEVICE_TYPE,
105
106
107
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
108
    RAG_TEMPLATE,
109
    ENABLE_RAG_LOCAL_WEB_FETCH,
110
    YOUTUBE_LOADER_LANGUAGE,
Timothy J. Baek's avatar
Timothy J. Baek committed
111
    ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
112
    RAG_WEB_SEARCH_ENGINE,
Timothy J. Baek's avatar
Timothy J. Baek committed
113
114
115
    SEARXNG_QUERY_URL,
    GOOGLE_PSE_API_KEY,
    GOOGLE_PSE_ENGINE_ID,
Timothy J. Baek's avatar
Timothy J. Baek committed
116
    BRAVE_SEARCH_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
117
118
119
    SERPSTACK_API_KEY,
    SERPSTACK_HTTPS,
    SERPER_API_KEY,
120
    SERPLY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
121
    RAG_WEB_SEARCH_RESULT_COUNT,
122
    RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
123
    RAG_EMBEDDING_OPENAI_BATCH_SIZE,
124
)
125

126
127
from constants import ERROR_MESSAGES

128
129
130
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
131
132
app = FastAPI()

133
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
134

135
136
137
138
139
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 = (
140
141
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
Steven Kreitzer's avatar
Steven Kreitzer committed
142

143
144
app.state.config.CHUNK_SIZE = CHUNK_SIZE
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
145

146
147
app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
148
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = RAG_EMBEDDING_OPENAI_BATCH_SIZE
149
150
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
151

152

153
154
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
155

156
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
157

Steven Kreitzer's avatar
Steven Kreitzer committed
158

159
app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
160
161
162
app.state.YOUTUBE_LOADER_TRANSLATION = None


Timothy J. Baek's avatar
Timothy J. Baek committed
163
app.state.config.ENABLE_RAG_WEB_SEARCH = ENABLE_RAG_WEB_SEARCH
Timothy J. Baek's avatar
Timothy J. Baek committed
164
165
app.state.config.RAG_WEB_SEARCH_ENGINE = RAG_WEB_SEARCH_ENGINE

Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
168
app.state.config.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
Timothy J. Baek's avatar
Timothy J. Baek committed
169
app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
170
171
172
app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
app.state.config.SERPER_API_KEY = SERPER_API_KEY
173
app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
174
175
176
177
app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = RAG_WEB_SEARCH_RESULT_COUNT
app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = RAG_WEB_SEARCH_CONCURRENT_REQUESTS


178
179
180
181
def update_embedding_model(
    embedding_model: str,
    update_model: bool = False,
):
182
    if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
        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(
207
    app.state.config.RAG_EMBEDDING_MODEL,
208
209
210
211
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)

update_reranking_model(
212
    app.state.config.RAG_RERANKING_MODEL,
213
214
    RAG_RERANKING_MODEL_AUTO_UPDATE,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
215

Timothy J. Baek's avatar
Timothy J. Baek committed
216
217

app.state.EMBEDDING_FUNCTION = get_embedding_function(
218
219
    app.state.config.RAG_EMBEDDING_ENGINE,
    app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
220
    app.state.sentence_transformer_ef,
221
222
    app.state.config.OPENAI_API_KEY,
    app.state.config.OPENAI_API_BASE_URL,
223
    app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
224
225
)

Timothy J. Baek's avatar
Timothy J. Baek committed
226
227
origins = ["*"]

228

Timothy J. Baek's avatar
Timothy J. Baek committed
229
230
231
232
233
234
235
236
237
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
238
class CollectionNameForm(BaseModel):
239
240
241
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
242
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
243
244
    url: str

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

246
247
248
249
class SearchForm(CollectionNameForm):
    query: str


Timothy J. Baek's avatar
Timothy J. Baek committed
250
251
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
252
253
    return {
        "status": True,
254
255
256
257
258
259
        "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,
260
        "openai_batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
261
262
263
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
264
265
@app.get("/embedding")
async def get_embedding_config(user=Depends(get_admin_user)):
266
267
    return {
        "status": True,
268
269
        "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
        "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
270
        "openai_config": {
271
272
            "url": app.state.config.OPENAI_API_BASE_URL,
            "key": app.state.config.OPENAI_API_KEY,
273
            "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
274
        },
275
276
277
    }


Steven Kreitzer's avatar
Steven Kreitzer committed
278
279
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
280
281
    return {
        "status": True,
282
        "reranking_model": app.state.config.RAG_RERANKING_MODEL,
283
    }
Steven Kreitzer's avatar
Steven Kreitzer committed
284
285


286
287
288
class OpenAIConfigForm(BaseModel):
    url: str
    key: str
289
    batch_size: Optional[int] = None
290
291


292
class EmbeddingModelUpdateForm(BaseModel):
293
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
294
    embedding_engine: str
295
296
297
    embedding_model: str


Timothy J. Baek's avatar
Timothy J. Baek committed
298
299
@app.post("/embedding/update")
async def update_embedding_config(
300
301
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
Self Denial's avatar
Self Denial committed
302
    log.info(
303
        f"Updating embedding model: {app.state.config.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
304
    )
305
    try:
306
307
        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
308

309
        if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
310
            if form_data.openai_config is not None:
311
312
                app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.config.OPENAI_API_KEY = form_data.openai_config.key
313
314
315
316
317
                app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = (
                    form_data.openai_config.batch_size
                    if form_data.openai_config.batch_size
                    else 1
                )
318

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

Timothy J. Baek's avatar
Timothy J. Baek committed
321
        app.state.EMBEDDING_FUNCTION = get_embedding_function(
322
323
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
324
            app.state.sentence_transformer_ef,
325
326
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
327
            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
328
329
        )

330
331
        return {
            "status": True,
332
333
            "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
            "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
334
            "openai_config": {
335
336
                "url": app.state.config.OPENAI_API_BASE_URL,
                "key": app.state.config.OPENAI_API_KEY,
337
                "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
338
            },
339
340
341
342
343
344
345
        }
    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
346
347


Steven Kreitzer's avatar
Steven Kreitzer committed
348
349
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
350

Steven Kreitzer's avatar
Steven Kreitzer committed
351
352
353
354
355
356

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

362
        update_reranking_model(app.state.config.RAG_RERANKING_MODEL), True
Steven Kreitzer's avatar
Steven Kreitzer committed
363
364
365

        return {
            "status": True,
366
            "reranking_model": app.state.config.RAG_RERANKING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
367
368
369
370
371
372
373
374
375
        }
    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
376
377
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
378
379
    return {
        "status": True,
380
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
381
        "chunk": {
382
383
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
384
        },
385
        "youtube": {
386
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
387
388
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
389
        "web": {
Timothy J. Baek's avatar
Timothy J. Baek committed
390
            "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Timothy J. Baek's avatar
Timothy J. Baek committed
391
            "search": {
Timothy J. Baek's avatar
Timothy J. Baek committed
392
                "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
393
                "engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
Timothy J. Baek's avatar
Timothy J. Baek committed
394
395
396
                "searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
                "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
                "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
Timothy J. Baek's avatar
Timothy J. Baek committed
397
                "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
398
399
400
                "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": app.state.config.SERPSTACK_HTTPS,
                "serper_api_key": app.state.config.SERPER_API_KEY,
401
                "serply_api_key": app.state.config.SERPLY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
402
403
                "result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
                "concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
Timothy J. Baek's avatar
Timothy J. Baek committed
404
            },
Timothy J. Baek's avatar
Timothy J. Baek committed
405
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
406
407
408
409
410
411
412
413
    }


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


414
415
416
417
418
class YoutubeLoaderConfig(BaseModel):
    language: List[str]
    translation: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
419
class WebSearchConfig(BaseModel):
Timothy J. Baek's avatar
Timothy J. Baek committed
420
    enabled: bool
Timothy J. Baek's avatar
Timothy J. Baek committed
421
    engine: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
422
423
424
    searxng_query_url: Optional[str] = None
    google_pse_api_key: Optional[str] = None
    google_pse_engine_id: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
425
    brave_search_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
426
427
428
    serpstack_api_key: Optional[str] = None
    serpstack_https: Optional[bool] = None
    serper_api_key: Optional[str] = None
429
    serply_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
430
431
432
433
    result_count: Optional[int] = None
    concurrent_requests: Optional[int] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
434
435
436
437
438
class WebConfig(BaseModel):
    search: WebSearchConfig
    web_loader_ssl_verification: Optional[bool] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
439
class ConfigUpdateForm(BaseModel):
440
441
    pdf_extract_images: Optional[bool] = None
    chunk: Optional[ChunkParamUpdateForm] = None
442
    youtube: Optional[YoutubeLoaderConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
443
    web: Optional[WebConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
444
445
446
447


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
448
    app.state.config.PDF_EXTRACT_IMAGES = (
449
        form_data.pdf_extract_images
450
451
        if form_data.pdf_extract_images is not None
        else app.state.config.PDF_EXTRACT_IMAGES
452
453
    )

Timothy J. Baek's avatar
Timothy J. Baek committed
454
455
456
    if form_data.chunk is not None:
        app.state.config.CHUNK_SIZE = form_data.chunk.chunk_size
        app.state.config.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
457

Timothy J. Baek's avatar
Timothy J. Baek committed
458
459
460
    if form_data.youtube is not None:
        app.state.config.YOUTUBE_LOADER_LANGUAGE = form_data.youtube.language
        app.state.YOUTUBE_LOADER_TRANSLATION = form_data.youtube.translation
461

Timothy J. Baek's avatar
Timothy J. Baek committed
462
463
464
465
    if form_data.web is not None:
        app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
            form_data.web.web_loader_ssl_verification
        )
466

Timothy J. Baek's avatar
Timothy J. Baek committed
467
        app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
Timothy J. Baek's avatar
Timothy J. Baek committed
468
469
470
471
472
473
474
475
476
477
478
479
        app.state.config.RAG_WEB_SEARCH_ENGINE = form_data.web.search.engine
        app.state.config.SEARXNG_QUERY_URL = form_data.web.search.searxng_query_url
        app.state.config.GOOGLE_PSE_API_KEY = form_data.web.search.google_pse_api_key
        app.state.config.GOOGLE_PSE_ENGINE_ID = (
            form_data.web.search.google_pse_engine_id
        )
        app.state.config.BRAVE_SEARCH_API_KEY = (
            form_data.web.search.brave_search_api_key
        )
        app.state.config.SERPSTACK_API_KEY = form_data.web.search.serpstack_api_key
        app.state.config.SERPSTACK_HTTPS = form_data.web.search.serpstack_https
        app.state.config.SERPER_API_KEY = form_data.web.search.serper_api_key
480
        app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
Timothy J. Baek's avatar
Timothy J. Baek committed
481
482
483
484
        app.state.config.RAG_WEB_SEARCH_RESULT_COUNT = form_data.web.search.result_count
        app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS = (
            form_data.web.search.concurrent_requests
        )
485

Timothy J. Baek's avatar
Timothy J. Baek committed
486
487
    return {
        "status": True,
488
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
489
        "chunk": {
490
491
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
492
        },
493
        "youtube": {
494
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
495
496
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
497
498
499
        "web": {
            "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
            "search": {
Timothy J. Baek's avatar
Timothy J. Baek committed
500
                "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
501
502
503
504
505
506
507
508
                "engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
                "searxng_query_url": app.state.config.SEARXNG_QUERY_URL,
                "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
                "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
                "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
                "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": app.state.config.SERPSTACK_HTTPS,
                "serper_api_key": app.state.config.SERPER_API_KEY,
509
                "serply_api_key": app.state.config.SERPLY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
510
511
512
513
                "result_count": app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
                "concurrent_requests": app.state.config.RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
            },
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
514
    }
515
516


Timothy J. Baek's avatar
Timothy J. Baek committed
517
518
519
520
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
521
        "template": app.state.config.RAG_TEMPLATE,
Timothy J. Baek's avatar
Timothy J. Baek committed
522
523
524
    }


525
526
527
528
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
529
530
531
532
        "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,
533
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
534
535


536
537
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
538
    r: Optional[float] = None
539
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
540
    hybrid: Optional[bool] = None
541
542
543
544
545
546


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
547
    app.state.config.RAG_TEMPLATE = (
Timothy J. Baek's avatar
Timothy J. Baek committed
548
        form_data.template if form_data.template else RAG_TEMPLATE
549
    )
550
551
552
    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
553
        form_data.hybrid if form_data.hybrid else False
554
    )
Steven Kreitzer's avatar
Steven Kreitzer committed
555
556
    return {
        "status": True,
557
558
559
560
        "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
561
    }
562
563


564
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
565
566
    collection_name: str
    query: str
567
    k: Optional[int] = None
568
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
569
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
570
571


572
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
573
def query_doc_handler(
574
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
575
576
    user=Depends(get_current_user),
):
577
    try:
578
        if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
Timothy J. Baek's avatar
Timothy J. Baek committed
579
580
581
            return query_doc_with_hybrid_search(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
582
                embedding_function=app.state.EMBEDDING_FUNCTION,
583
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
584
                reranking_function=app.state.sentence_transformer_rf,
585
                r=(
586
                    form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
587
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
588
589
590
591
592
            )
        else:
            return query_doc(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
593
                embedding_function=app.state.EMBEDDING_FUNCTION,
594
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
595
            )
596
    except Exception as e:
597
        log.exception(e)
598
599
600
601
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
602
603


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
604
605
606
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
607
    k: Optional[int] = None
608
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
609
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
610
611


612
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
613
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
614
615
616
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
617
    try:
618
        if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
Timothy J. Baek's avatar
Timothy J. Baek committed
619
620
621
            return query_collection_with_hybrid_search(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
622
                embedding_function=app.state.EMBEDDING_FUNCTION,
623
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
624
                reranking_function=app.state.sentence_transformer_rf,
625
                r=(
626
                    form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
627
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
628
629
630
631
632
            )
        else:
            return query_collection(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
633
                embedding_function=app.state.EMBEDDING_FUNCTION,
634
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
635
            )
636

637
638
639
640
641
642
    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
643
644


Timothy J. Baek's avatar
Timothy J. Baek committed
645
646
647
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
648
649
650
        loader = YoutubeLoader.from_youtube_url(
            form_data.url,
            add_video_info=True,
651
            language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
652
653
            translation=app.state.YOUTUBE_LOADER_TRANSLATION,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
        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),
        )


674
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
675
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
676
677
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
678
        loader = get_web_loader(
679
            form_data.url,
680
            verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
681
        )
682
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
683
684
685
686
687

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

688
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
689
690
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
691
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
692
693
            "filename": form_data.url,
        }
694
    except Exception as e:
695
        log.exception(e)
696
697
698
699
700
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )

701

702
def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
703
    # Check if the URL is valid
704
    if not validate_url(url):
705
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
706
    return SafeWebBaseLoader(
707
708
709
        url,
        verify_ssl=verify_ssl,
        requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
710
        continue_on_failure=True,
711
    )
712
713


714
715
716
717
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)
718
        if not ENABLE_RAG_LOCAL_WEB_FETCH:
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
            # 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


738
739
740
741
742
743
744
745
746
747
748
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


Timothy J. Baek's avatar
Timothy J. Baek committed
749
750
751
752
753
754
755
756
def search_web(engine: str, query: str) -> list[SearchResult]:
    """Search the web using a search engine and return the results as a list of SearchResult objects.
    Will look for a search engine API key in environment variables in the following order:
    - SEARXNG_QUERY_URL
    - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
    - BRAVE_SEARCH_API_KEY
    - SERPSTACK_API_KEY
    - SERPER_API_KEY
757
    - SERPLY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
758
759
760
761
762
763
764
765

    Args:
        query (str): The query to search for
    """

    # TODO: add playwright to search the web
    if engine == "searxng":
        if app.state.config.SEARXNG_QUERY_URL:
Timothy J. Baek's avatar
Timothy J. Baek committed
766
767
768
769
770
            return search_searxng(
                app.state.config.SEARXNG_QUERY_URL,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
771
772
773
774
775
776
777
778
779
780
781
        else:
            raise Exception("No SEARXNG_QUERY_URL found in environment variables")
    elif engine == "google_pse":
        if (
            app.state.config.GOOGLE_PSE_API_KEY
            and app.state.config.GOOGLE_PSE_ENGINE_ID
        ):
            return search_google_pse(
                app.state.config.GOOGLE_PSE_API_KEY,
                app.state.config.GOOGLE_PSE_ENGINE_ID,
                query,
Timothy J. Baek's avatar
Timothy J. Baek committed
782
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
Timothy J. Baek's avatar
Timothy J. Baek committed
783
784
785
786
787
788
789
            )
        else:
            raise Exception(
                "No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
            )
    elif engine == "brave":
        if app.state.config.BRAVE_SEARCH_API_KEY:
Timothy J. Baek's avatar
Timothy J. Baek committed
790
791
792
793
794
            return search_brave(
                app.state.config.BRAVE_SEARCH_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
795
796
797
798
799
800
801
        else:
            raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
    elif engine == "serpstack":
        if app.state.config.SERPSTACK_API_KEY:
            return search_serpstack(
                app.state.config.SERPSTACK_API_KEY,
                query,
Timothy J. Baek's avatar
Timothy J. Baek committed
802
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
Timothy J. Baek's avatar
Timothy J. Baek committed
803
804
805
806
807
808
                https_enabled=app.state.config.SERPSTACK_HTTPS,
            )
        else:
            raise Exception("No SERPSTACK_API_KEY found in environment variables")
    elif engine == "serper":
        if app.state.config.SERPER_API_KEY:
Timothy J. Baek's avatar
Timothy J. Baek committed
809
810
811
812
813
            return search_serper(
                app.state.config.SERPER_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
814
815
        else:
            raise Exception("No SERPER_API_KEY found in environment variables")
816
817
818
819
820
821
822
823
824
    elif engine == "serply":
        if app.state.config.SERPLY_API_KEY:
            return search_serply(
                app.state.config.SERPLY_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
            )
        else:
            raise Exception("No SERPLY_API_KEY found in environment variables")
825
826
    elif engine == "duckduckgo":
        return search_duckduckgo(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT)
Timothy J. Baek's avatar
Timothy J. Baek committed
827
828
829
830
    else:
        raise Exception("No search engine API key found in environment variables")


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
831
832
@app.post("/web/search")
def store_web_search(form_data: SearchForm, user=Depends(get_current_user)):
833
    try:
834
        logging.info(f"trying to web search with {app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query}")
Timothy J. Baek's avatar
Timothy J. Baek committed
835
836
837
838
839
840
841
842
843
844
845
846
847
        web_results = search_web(
            app.state.config.RAG_WEB_SEARCH_ENGINE, form_data.query
        )
    except Exception as e:
        log.exception(e)

        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.WEB_SEARCH_ERROR(e),
        )

    try:
848
849
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
850
        data = loader.load()
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869

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


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

872
    text_splitter = RecursiveCharacterTextSplitter(
873
874
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
875
876
        add_start_index=True,
    )
877

878
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
879
880

    if len(docs) > 0:
881
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
882
883
884
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
885
886
887


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
888
    text, metadata, collection_name, overwrite: bool = False
889
890
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
891
892
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
893
894
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
895
    docs = text_splitter.create_documents([text], metadatas=[metadata])
896
897
898
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

902
903
904
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]

mindspawn's avatar
mindspawn committed
905
906
907
908
909
910
911
    # ChromaDB does not like datetime formats
    # for meta-data so convert them to string.
    for metadata in metadatas:
        for key, value in metadata.items():
            if isinstance(value, datetime):
                metadata[key] = str(value)

912
913
914
915
    try:
        if overwrite:
            for collection in CHROMA_CLIENT.list_collections():
                if collection_name == collection.name:
916
                    log.info(f"deleting existing collection {collection_name}")
917
918
                    CHROMA_CLIENT.delete_collection(name=collection_name)

919
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
920

Timothy J. Baek's avatar
Timothy J. Baek committed
921
        embedding_func = get_embedding_function(
922
923
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
924
            app.state.sentence_transformer_ef,
925
926
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
927
            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Steven Kreitzer's avatar
Steven Kreitzer committed
928
929
930
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
931
        embeddings = embedding_func(embedding_texts)
932
933
934

        for batch in create_batches(
            api=CHROMA_CLIENT,
935
            ids=[str(uuid.uuid4()) for _ in texts],
936
937
938
939
940
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
941

942
        return True
943
    except Exception as e:
944
        log.exception(e)
945
946
947
948
949
950
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


951
952
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
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
    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",
mindspawn's avatar
mindspawn committed
998
        "msg",
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
999
1000
1001
    ]

    if file_ext == "pdf":
1002
        loader = PyPDFLoader(
1003
            file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES
1004
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1005
1006
1007
1008
1009
1010
    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)
1011
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
1012
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1013
1014
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
1015
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1016
1017
        loader = UnstructuredEPubLoader(file_path)
    elif (
1018
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1019
1020
1021
1022
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
1023
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1024
1025
1026
1027
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
1028
1029
1030
1031
1032
    elif file_content_type in [
        "application/vnd.ms-powerpoint",
        "application/vnd.openxmlformats-officedocument.presentationml.presentation",
    ] or file_ext in ["ppt", "pptx"]:
        loader = UnstructuredPowerPointLoader(file_path)
mindspawn's avatar
mindspawn committed
1033
1034
    elif file_ext == "msg":
        loader = OutlookMessageLoader(file_path)
1035
1036
1037
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
1038
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1039
    else:
1040
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1041
1042
1043
1044
1045
        known_type = False

    return loader, known_type


1046
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
1047
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
1048
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
1049
1050
1051
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
1052
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
1053

1054
    log.info(f"file.content_type: {file.content_type}")
1055
    try:
1056
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
1057
        filename = os.path.basename(unsanitized_filename)
1058

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

1061
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
1062
        with open(file_path, "wb") as f:
1063
1064
1065
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
1066
1067
1068
1069
1070
        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
1071
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
1072
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084

        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
1085
1086
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
1087
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
1088
            )
1089
    except Exception as e:
1090
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
        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),
            )
1101
1102


1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
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
1119
1120
1121
1122
1123
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133

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


1134
1135
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
1136
1137
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
1138
1139
1140
1141
1142
1143
1144
1145
1146
            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
1147
1148
1149
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
1150
1151
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
                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,
                                                        )
1176
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
1177
1178
1179
1180
1181
1182
1183
1184
1185
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
1186
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
1187
                    pass
1188

1189
        except Exception as e:
1190
            log.exception(e)
1191
1192
1193
1194

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
1195
@app.get("/reset/db")
1196
1197
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
1198
1199


Timothy J. Baek's avatar
Timothy J. Baek committed
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
@app.get("/reset/uploads")
def reset_upload_dir(user=Depends(get_admin_user)) -> bool:
    folder = f"{UPLOAD_DIR}"
    try:
        # Check if the directory exists
        if os.path.exists(folder):
            # Iterate over all the files and directories in the specified directory
            for filename in os.listdir(folder):
                file_path = os.path.join(folder, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)  # Remove the file or link
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)  # Remove the directory
                except Exception as e:
                    print(f"Failed to delete {file_path}. Reason: {e}")
        else:
            print(f"The directory {folder} does not exist")
    except Exception as e:
        print(f"Failed to process the directory {folder}. Reason: {e}")

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
1224
@app.get("/reset")
1225
1226
1227
1228
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
1229
        try:
1230
1231
1232
1233
            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
1234
        except Exception as e:
1235
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
1236

1237
1238
1239
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
1240
        log.exception(e)
1241
1242

    return True
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1243

1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
class SafeWebBaseLoader(WebBaseLoader):
    """WebBaseLoader with enhanced error handling for URLs."""
    def lazy_load(self) -> Iterator[Document]:
        """Lazy load text from the url(s) in web_path with error handling."""
        for path in self.web_paths:
            try:
                soup = self._scrape(path, bs_kwargs=self.bs_kwargs)
                text = soup.get_text(**self.bs_get_text_kwargs)

                # Build metadata
                metadata = {"source": path}
                if title := soup.find("title"):
                    metadata["title"] = title.get_text()
                if description := soup.find("meta", attrs={"name": "description"}):
                    metadata["description"] = description.get("content", "No description found.")
                if html := soup.find("html"):
                    metadata["language"] = html.get("lang", "No language found.")
                
                yield Document(page_content=text, metadata=metadata)
            except Exception as e:
                # Log the error and continue with the next URL
                log.error(f"Error loading {path}: {e}")
                
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1267
1268
1269
1270
1271
1272
1273
1274
1275
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)}