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

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

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

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

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


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

51
import sentence_transformers
52

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
68
69
70
71
72
73
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
74
from apps.rag.search.serply import search_serply
75
from apps.rag.search.duckduckgo import search_duckduckgo
76
from apps.rag.search.tavily import search_tavily
Timothy J. Baek's avatar
Timothy J. Baek committed
77

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

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

130
131
from constants import ERROR_MESSAGES

132
133
134
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
135
136
app = FastAPI()

137
app.state.config = AppConfig()
Timothy J. Baek's avatar
Timothy J. Baek committed
138

139
140
141
142
143
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 = (
144
145
    ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
Steven Kreitzer's avatar
Steven Kreitzer committed
146

147
148
app.state.config.CHUNK_SIZE = CHUNK_SIZE
app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
149

150
151
app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
152
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = RAG_EMBEDDING_OPENAI_BATCH_SIZE
153
154
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
155

156

157
158
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
159

160
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
161

Steven Kreitzer's avatar
Steven Kreitzer committed
162

163
app.state.config.YOUTUBE_LOADER_LANGUAGE = YOUTUBE_LOADER_LANGUAGE
164
165
166
app.state.YOUTUBE_LOADER_TRANSLATION = None


Timothy J. Baek's avatar
Timothy J. Baek committed
167
app.state.config.ENABLE_RAG_WEB_SEARCH = ENABLE_RAG_WEB_SEARCH
Timothy J. Baek's avatar
Timothy J. Baek committed
168
app.state.config.RAG_WEB_SEARCH_ENGINE = RAG_WEB_SEARCH_ENGINE
169
app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS = RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
Timothy J. Baek's avatar
Timothy J. Baek committed
170

Timothy J. Baek's avatar
Timothy J. Baek committed
171
172
173
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
174
app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
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
178
app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
179
app.state.config.TAVILY_API_KEY = TAVILY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
180
181
182
183
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


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

update_reranking_model(
218
    app.state.config.RAG_RERANKING_MODEL,
219
220
    RAG_RERANKING_MODEL_AUTO_UPDATE,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
221

Timothy J. Baek's avatar
Timothy J. Baek committed
222
223

app.state.EMBEDDING_FUNCTION = get_embedding_function(
224
225
    app.state.config.RAG_EMBEDDING_ENGINE,
    app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
226
    app.state.sentence_transformer_ef,
227
228
    app.state.config.OPENAI_API_KEY,
    app.state.config.OPENAI_API_BASE_URL,
229
    app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
230
231
)

Timothy J. Baek's avatar
Timothy J. Baek committed
232
233
origins = ["*"]

234

Timothy J. Baek's avatar
Timothy J. Baek committed
235
236
237
238
239
240
241
242
243
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
244
class CollectionNameForm(BaseModel):
245
246
247
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
248
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
249
250
    url: str

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

252
253
254
255
class SearchForm(CollectionNameForm):
    query: str


Timothy J. Baek's avatar
Timothy J. Baek committed
256
257
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
258
259
    return {
        "status": True,
260
261
262
263
264
265
        "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,
266
        "openai_batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
267
268
269
    }


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


Steven Kreitzer's avatar
Steven Kreitzer committed
284
285
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
286
287
    return {
        "status": True,
288
        "reranking_model": app.state.config.RAG_RERANKING_MODEL,
289
    }
Steven Kreitzer's avatar
Steven Kreitzer committed
290
291


292
293
294
class OpenAIConfigForm(BaseModel):
    url: str
    key: str
295
    batch_size: Optional[int] = None
296
297


298
class EmbeddingModelUpdateForm(BaseModel):
299
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
300
    embedding_engine: str
301
302
303
    embedding_model: str


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

315
        if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
316
            if form_data.openai_config is not None:
317
318
                app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.config.OPENAI_API_KEY = form_data.openai_config.key
319
320
321
322
323
                app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = (
                    form_data.openai_config.batch_size
                    if form_data.openai_config.batch_size
                    else 1
                )
324

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

Timothy J. Baek's avatar
Timothy J. Baek committed
327
        app.state.EMBEDDING_FUNCTION = get_embedding_function(
328
329
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Timothy J. Baek's avatar
Timothy J. Baek committed
330
            app.state.sentence_transformer_ef,
331
332
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
333
            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Timothy J. Baek's avatar
Timothy J. Baek committed
334
335
        )

336
337
        return {
            "status": True,
338
339
            "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
            "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
340
            "openai_config": {
341
342
                "url": app.state.config.OPENAI_API_BASE_URL,
                "key": app.state.config.OPENAI_API_KEY,
343
                "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
344
            },
345
346
347
348
349
350
351
        }
    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
352
353


Steven Kreitzer's avatar
Steven Kreitzer committed
354
355
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
356

Steven Kreitzer's avatar
Steven Kreitzer committed
357
358
359
360
361
362

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

368
        update_reranking_model(app.state.config.RAG_RERANKING_MODEL), True
Steven Kreitzer's avatar
Steven Kreitzer committed
369
370
371

        return {
            "status": True,
372
            "reranking_model": app.state.config.RAG_RERANKING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
373
374
375
376
377
378
379
380
381
        }
    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
382
383
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
384
385
    return {
        "status": True,
386
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
387
        "chunk": {
388
389
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
390
        },
391
        "youtube": {
392
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
393
394
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
395
        "web": {
Timothy J. Baek's avatar
Timothy J. Baek committed
396
            "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
Timothy J. Baek's avatar
Timothy J. Baek committed
397
            "search": {
Timothy J. Baek's avatar
Timothy J. Baek committed
398
                "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
399
                "engine": app.state.config.RAG_WEB_SEARCH_ENGINE,
Timothy J. Baek's avatar
Timothy J. Baek committed
400
401
402
                "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
403
                "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
404
405
406
                "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,
407
                "serply_api_key": app.state.config.SERPLY_API_KEY,
408
                "tavily_api_key": app.state.config.TAVILY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
409
410
                "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
411
            },
Timothy J. Baek's avatar
Timothy J. Baek committed
412
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
413
414
415
416
417
418
419
420
    }


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


421
422
423
424
425
class YoutubeLoaderConfig(BaseModel):
    language: List[str]
    translation: Optional[str] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
426
class WebSearchConfig(BaseModel):
Timothy J. Baek's avatar
Timothy J. Baek committed
427
    enabled: bool
Timothy J. Baek's avatar
Timothy J. Baek committed
428
    engine: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
429
430
431
    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
432
    brave_search_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
433
434
435
    serpstack_api_key: Optional[str] = None
    serpstack_https: Optional[bool] = None
    serper_api_key: Optional[str] = None
436
    serply_api_key: Optional[str] = None
437
    tavily_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
438
439
440
441
    result_count: Optional[int] = None
    concurrent_requests: Optional[int] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
442
443
444
445
446
class WebConfig(BaseModel):
    search: WebSearchConfig
    web_loader_ssl_verification: Optional[bool] = None


Timothy J. Baek's avatar
Timothy J. Baek committed
447
class ConfigUpdateForm(BaseModel):
448
449
    pdf_extract_images: Optional[bool] = None
    chunk: Optional[ChunkParamUpdateForm] = None
450
    youtube: Optional[YoutubeLoaderConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
451
    web: Optional[WebConfig] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
452
453
454
455


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
456
    app.state.config.PDF_EXTRACT_IMAGES = (
457
        form_data.pdf_extract_images
458
459
        if form_data.pdf_extract_images is not None
        else app.state.config.PDF_EXTRACT_IMAGES
460
461
    )

Timothy J. Baek's avatar
Timothy J. Baek committed
462
463
464
    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
465

Timothy J. Baek's avatar
Timothy J. Baek committed
466
467
468
    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
469

Timothy J. Baek's avatar
Timothy J. Baek committed
470
471
472
473
    if form_data.web is not None:
        app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
            form_data.web.web_loader_ssl_verification
        )
474

Timothy J. Baek's avatar
Timothy J. Baek committed
475
        app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
Timothy J. Baek's avatar
Timothy J. Baek committed
476
477
478
479
480
481
482
483
484
485
486
487
        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
488
        app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
489
        app.state.config.TAVILY_API_KEY = form_data.web.search.tavily_api_key
Timothy J. Baek's avatar
Timothy J. Baek committed
490
491
492
493
        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
        )
494

Timothy J. Baek's avatar
Timothy J. Baek committed
495
496
    return {
        "status": True,
497
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
498
        "chunk": {
499
500
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
501
        },
502
        "youtube": {
503
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
504
505
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
506
507
508
        "web": {
            "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
            "search": {
Timothy J. Baek's avatar
Timothy J. Baek committed
509
                "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
510
511
512
513
514
515
516
517
                "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,
518
                "serply_api_key": app.state.config.SERPLY_API_KEY,
519
                "tavily_api_key": app.state.config.TAVILY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
520
521
522
523
                "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
524
    }
525
526


Timothy J. Baek's avatar
Timothy J. Baek committed
527
528
529
530
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
531
        "template": app.state.config.RAG_TEMPLATE,
Timothy J. Baek's avatar
Timothy J. Baek committed
532
533
534
    }


535
536
537
538
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
539
540
541
542
        "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,
543
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
544
545


546
547
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
548
    r: Optional[float] = None
549
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
550
    hybrid: Optional[bool] = None
551
552
553
554
555
556


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
557
    app.state.config.RAG_TEMPLATE = (
Timothy J. Baek's avatar
Timothy J. Baek committed
558
        form_data.template if form_data.template else RAG_TEMPLATE
559
    )
560
561
562
    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
563
        form_data.hybrid if form_data.hybrid else False
564
    )
Steven Kreitzer's avatar
Steven Kreitzer committed
565
566
    return {
        "status": True,
567
568
569
570
        "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
571
    }
572
573


574
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
575
576
    collection_name: str
    query: str
577
    k: Optional[int] = None
578
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
579
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
580
581


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


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
614
615
616
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
617
    k: Optional[int] = None
618
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
619
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
620
621


622
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
623
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
624
625
626
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
627
    try:
628
        if app.state.config.ENABLE_RAG_HYBRID_SEARCH:
Timothy J. Baek's avatar
Timothy J. Baek committed
629
630
631
            return query_collection_with_hybrid_search(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
632
                embedding_function=app.state.EMBEDDING_FUNCTION,
633
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
634
                reranking_function=app.state.sentence_transformer_rf,
635
                r=(
636
                    form_data.r if form_data.r else app.state.config.RELEVANCE_THRESHOLD
637
                ),
Timothy J. Baek's avatar
Timothy J. Baek committed
638
639
640
641
642
            )
        else:
            return query_collection(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
643
                embedding_function=app.state.EMBEDDING_FUNCTION,
644
                k=form_data.k if form_data.k else app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
645
            )
646

647
648
649
650
651
652
    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
653
654


Timothy J. Baek's avatar
Timothy J. Baek committed
655
656
657
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
658
659
660
        loader = YoutubeLoader.from_youtube_url(
            form_data.url,
            add_video_info=True,
661
            language=app.state.config.YOUTUBE_LOADER_LANGUAGE,
662
663
            translation=app.state.YOUTUBE_LOADER_TRANSLATION,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
        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),
        )


684
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
685
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
686
687
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
688
        loader = get_web_loader(
689
            form_data.url,
690
            verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
691
        )
692
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
693
694
695
696
697

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

698
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
699
700
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
701
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
702
703
            "filename": form_data.url,
        }
704
    except Exception as e:
705
        log.exception(e)
706
707
708
709
710
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )

711

712
def get_web_loader(url: Union[str, Sequence[str]], verify_ssl: bool = True):
713
    # Check if the URL is valid
714
    if not validate_url(url):
715
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
716
    return SafeWebBaseLoader(
717
718
719
        url,
        verify_ssl=verify_ssl,
        requests_per_second=RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
720
        continue_on_failure=True,
721
    )
722
723


724
725
726
727
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)
728
        if not ENABLE_RAG_LOCAL_WEB_FETCH:
Timothy J. Baek's avatar
revert  
Timothy J. Baek committed
729
730
731
732
733
734
735
736
737
738
739
            # 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):
740
741
742
743
744
745
746
                    raise ValueError(ERROR_MESSAGES.INVALID_URL)
        return True
    elif isinstance(url, Sequence):
        return all(validate_url(u) for u in url)
    else:
        return False

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

Timothy J. Baek's avatar
revert  
Timothy J. Baek committed
748
749
750
751
752
753
754
755
756
757
758
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
759
760
761
762
763
764
765
766
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
767
    - SERPLY_API_KEY
768
    - TAVILY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
769
770
771
772
773
774
775
    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
776
777
778
779
            return search_searxng(
                app.state.config.SEARXNG_QUERY_URL,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
780
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
Timothy J. Baek's avatar
Timothy J. Baek committed
781
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
782
783
784
785
786
787
788
789
790
791
792
        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
793
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
794
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
Timothy J. Baek's avatar
Timothy J. Baek committed
795
796
797
798
799
800
801
            )
        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
802
803
804
805
            return search_brave(
                app.state.config.BRAVE_SEARCH_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
806
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
Timothy J. Baek's avatar
Timothy J. Baek committed
807
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
808
809
810
811
812
813
814
        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
815
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
816
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS,
Timothy J. Baek's avatar
Timothy J. Baek committed
817
818
819
820
821
822
                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
823
824
825
826
            return search_serper(
                app.state.config.SERPER_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
827
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
Timothy J. Baek's avatar
Timothy J. Baek committed
828
            )
Timothy J. Baek's avatar
Timothy J. Baek committed
829
830
        else:
            raise Exception("No SERPER_API_KEY found in environment variables")
831
832
833
834
835
836
    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,
837
                app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS
838
839
840
            )
        else:
            raise Exception("No SERPLY_API_KEY found in environment variables")
841
    elif engine == "duckduckgo":
842
        return search_duckduckgo(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT, app.state.config.RAG_WEB_SEARCH_WHITE_LIST_DOMAINS)
843
844
845
846
847
848
849
850
851
    elif engine == "tavily":
        if app.state.config.TAVILY_API_KEY:
            return search_tavily(
                app.state.config.TAVILY_API_KEY,
                query,
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
            )
        else:
            raise Exception("No TAVILY_API_KEY found in environment variables")
Timothy J. Baek's avatar
Timothy J. Baek committed
852
853
854
855
    else:
        raise Exception("No search engine API key found in environment variables")


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
856
857
@app.post("/web/search")
def store_web_search(form_data: SearchForm, user=Depends(get_current_user)):
858
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
859
860
861
        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
862
863
864
865
866
867
868
869
870
871
872
873
874
        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:
875
876
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
877
        data = loader.load()
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896

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


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

899
    text_splitter = RecursiveCharacterTextSplitter(
900
901
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
902
903
        add_start_index=True,
    )
904

905
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
906
907

    if len(docs) > 0:
908
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
909
910
911
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
912
913
914


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
915
    text, metadata, collection_name, overwrite: bool = False
916
917
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
918
919
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
920
921
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
922
    docs = text_splitter.create_documents([text], metadatas=[metadata])
923
924
925
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

929
930
931
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]

mindspawn's avatar
mindspawn committed
932
933
934
935
936
937
938
    # 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)

939
940
941
942
    try:
        if overwrite:
            for collection in CHROMA_CLIENT.list_collections():
                if collection_name == collection.name:
943
                    log.info(f"deleting existing collection {collection_name}")
944
945
                    CHROMA_CLIENT.delete_collection(name=collection_name)

946
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
947

Timothy J. Baek's avatar
Timothy J. Baek committed
948
        embedding_func = get_embedding_function(
949
950
            app.state.config.RAG_EMBEDDING_ENGINE,
            app.state.config.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
951
            app.state.sentence_transformer_ef,
952
953
            app.state.config.OPENAI_API_KEY,
            app.state.config.OPENAI_API_BASE_URL,
954
            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
Steven Kreitzer's avatar
Steven Kreitzer committed
955
956
957
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
958
        embeddings = embedding_func(embedding_texts)
959
960
961

        for batch in create_batches(
            api=CHROMA_CLIENT,
962
            ids=[str(uuid.uuid4()) for _ in texts],
963
964
965
966
967
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
968

969
        return True
970
    except Exception as e:
971
        log.exception(e)
972
973
974
975
976
977
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


978
979
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
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
    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
1025
        "msg",
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1026
1027
1028
    ]

    if file_ext == "pdf":
1029
        loader = PyPDFLoader(
1030
            file_path, extract_images=app.state.config.PDF_EXTRACT_IMAGES
1031
        )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1032
1033
1034
1035
1036
1037
    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)
1038
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
1039
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1040
1041
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
1042
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1043
1044
        loader = UnstructuredEPubLoader(file_path)
    elif (
1045
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1046
1047
1048
1049
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
1050
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1051
1052
1053
1054
        "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
1055
1056
1057
1058
1059
    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
1060
1061
    elif file_ext == "msg":
        loader = OutlookMessageLoader(file_path)
1062
1063
1064
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
1065
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1066
    else:
1067
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1068
1069
1070
1071
1072
        known_type = False

    return loader, known_type


1073
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
1074
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
1075
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
1076
1077
1078
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
1079
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
1080

1081
    log.info(f"file.content_type: {file.content_type}")
1082
    try:
1083
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
1084
        filename = os.path.basename(unsanitized_filename)
1085

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

1088
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
1089
        with open(file_path, "wb") as f:
1090
1091
1092
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
1093
1094
1095
1096
1097
        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
1098
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
1099
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111

        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
1112
1113
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
1114
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
1115
            )
1116
    except Exception as e:
1117
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
        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),
            )
1128
1129


1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
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
1146
1147
1148
1149
1150
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160

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


1161
1162
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
1163
1164
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
1165
1166
1167
1168
1169
1170
1171
1172
1173
            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
1174
1175
1176
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
1177
1178
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
                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,
                                                        )
1203
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
1204
1205
1206
1207
1208
1209
1210
1211
1212
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
1213
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
1214
                    pass
1215

1216
        except Exception as e:
1217
            log.exception(e)
1218
1219
1220
1221

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
1222
@app.get("/reset/db")
1223
1224
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
1225
1226


Timothy J. Baek's avatar
Timothy J. Baek committed
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
@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
1251
@app.get("/reset")
1252
1253
1254
1255
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
1256
        try:
1257
1258
1259
1260
            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
1261
        except Exception as e:
1262
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
1263

1264
1265
1266
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
1267
        log.exception(e)
1268
1269

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

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

1272
1273
class SafeWebBaseLoader(WebBaseLoader):
    """WebBaseLoader with enhanced error handling for URLs."""
Timothy J. Baek's avatar
Timothy J. Baek committed
1274

1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
    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"}):
Timothy J. Baek's avatar
Timothy J. Baek committed
1287
1288
1289
                    metadata["description"] = description.get(
                        "content", "No description found."
                    )
1290
1291
                if html := soup.find("html"):
                    metadata["language"] = html.get("lang", "No language found.")
Timothy J. Baek's avatar
Timothy J. Baek committed
1292

1293
1294
1295
1296
                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
Timothy J. Baek committed
1297
1298


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1299
1300
1301
1302
1303
1304
1305
1306
1307
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)}