"vscode:/vscode.git/clone" did not exist on "9dc0b779a25a929cd59e1f914abf6aa9c33b278f"
main.py 41.9 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
Timothy J. Baek's avatar
Timothy J. Baek committed
76

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

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

127
128
from constants import ERROR_MESSAGES

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

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

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

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

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

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

153

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

157
app.state.config.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
158

Steven Kreitzer's avatar
Steven Kreitzer committed
159

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


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

Timothy J. Baek's avatar
Timothy J. Baek committed
167
168
169
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
170
app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
171
172
173
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
174
app.state.config.SERPLY_API_KEY = SERPLY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
178
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


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

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

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

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

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

229

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


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


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

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

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


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


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


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


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


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


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

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

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

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

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


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

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

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

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

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


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


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


Timothy J. Baek's avatar
Timothy J. Baek committed
420
class WebSearchConfig(BaseModel):
Timothy J. Baek's avatar
Timothy J. Baek committed
421
    enabled: bool
Timothy J. Baek's avatar
Timothy J. Baek committed
422
    engine: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
423
424
425
    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
426
    brave_search_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
427
428
429
    serpstack_api_key: Optional[str] = None
    serpstack_https: Optional[bool] = None
    serper_api_key: Optional[str] = None
430
    serply_api_key: Optional[str] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
431
432
433
434
    result_count: Optional[int] = None
    concurrent_requests: Optional[int] = None


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


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


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

Timothy J. Baek's avatar
Timothy J. Baek committed
455
456
457
    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
458

Timothy J. Baek's avatar
Timothy J. Baek committed
459
460
461
    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
462

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

Timothy J. Baek's avatar
Timothy J. Baek committed
468
        app.state.config.ENABLE_RAG_WEB_SEARCH = form_data.web.search.enabled
Timothy J. Baek's avatar
Timothy J. Baek committed
469
470
471
472
473
474
475
476
477
478
479
480
        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
481
        app.state.config.SERPLY_API_KEY = form_data.web.search.serply_api_key
Timothy J. Baek's avatar
Timothy J. Baek committed
482
483
484
485
        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
        )
486

Timothy J. Baek's avatar
Timothy J. Baek committed
487
488
    return {
        "status": True,
489
        "pdf_extract_images": app.state.config.PDF_EXTRACT_IMAGES,
Timothy J. Baek's avatar
Timothy J. Baek committed
490
        "chunk": {
491
492
            "chunk_size": app.state.config.CHUNK_SIZE,
            "chunk_overlap": app.state.config.CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
493
        },
494
        "youtube": {
495
            "language": app.state.config.YOUTUBE_LOADER_LANGUAGE,
496
497
            "translation": app.state.YOUTUBE_LOADER_TRANSLATION,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
498
499
500
        "web": {
            "ssl_verification": app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
            "search": {
Timothy J. Baek's avatar
Timothy J. Baek committed
501
                "enabled": app.state.config.ENABLE_RAG_WEB_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
502
503
504
505
506
507
508
509
                "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,
510
                "serply_api_key": app.state.config.SERPLY_API_KEY,
Timothy J. Baek's avatar
Timothy J. Baek committed
511
512
513
514
                "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
515
    }
516
517


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


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


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


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


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


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


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


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

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


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


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

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

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

702

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


715
716
717
718
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)
719
        if not ENABLE_RAG_LOCAL_WEB_FETCH:
Timothy J. Baek's avatar
revert  
Timothy J. Baek committed
720
721
722
723
724
725
726
727
728
729
730
            # 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):
731
732
733
734
735
736
737
                    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
738

Timothy J. Baek's avatar
revert  
Timothy J. Baek committed
739
740
741
742
743
744
745
746
747
748
749
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
750
751
752
753
754
755
756
757
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
758
    - SERPLY_API_KEY
Timothy J. Baek's avatar
Timothy J. Baek committed
759
760
761
762
763
764
765
766

    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
767
768
769
770
771
            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
772
773
774
775
776
777
778
779
780
781
782
        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
783
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
Timothy J. Baek's avatar
Timothy J. Baek committed
784
785
786
787
788
789
790
            )
        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
791
792
793
794
795
            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
796
797
798
799
800
801
802
        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
803
                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
Timothy J. Baek's avatar
Timothy J. Baek committed
804
805
806
807
808
809
                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
810
811
812
813
814
            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
815
816
        else:
            raise Exception("No SERPER_API_KEY found in environment variables")
817
818
819
820
821
822
823
824
825
    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")
826
827
    elif engine == "duckduckgo":
        return search_duckduckgo(query, app.state.config.RAG_WEB_SEARCH_RESULT_COUNT)
Timothy J. Baek's avatar
Timothy J. Baek committed
828
829
830
831
    else:
        raise Exception("No search engine API key found in environment variables")


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
832
833
@app.post("/web/search")
def store_web_search(form_data: SearchForm, user=Depends(get_current_user)):
834
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
835
836
837
        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
838
839
840
841
842
843
844
845
846
847
848
849
850
        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:
851
852
        urls = [result.link for result in web_results]
        loader = get_web_loader(urls)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
853
        data = loader.load()
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872

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


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

875
    text_splitter = RecursiveCharacterTextSplitter(
876
877
        chunk_size=app.state.config.CHUNK_SIZE,
        chunk_overlap=app.state.config.CHUNK_OVERLAP,
878
879
        add_start_index=True,
    )
880

881
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
882
883

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


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


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

905
906
907
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]

mindspawn's avatar
mindspawn committed
908
909
910
911
912
913
914
    # 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)

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

922
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
923

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

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
934
        embeddings = embedding_func(embedding_texts)
935
936
937

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

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

        return False


954
955
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
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
998
999
1000
    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
1001
        "msg",
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1002
1003
1004
    ]

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

    return loader, known_type


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

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

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

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

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

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


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

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


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

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

1192
        except Exception as e:
1193
            log.exception(e)
1194
1195
1196
1197

    return True


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


Timothy J. Baek's avatar
Timothy J. Baek committed
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
@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
1227
@app.get("/reset")
1228
1229
1230
1231
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
1232
        try:
1233
1234
1235
1236
            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
1237
        except Exception as e:
1238
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
1239

1240
1241
1242
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
1243
        log.exception(e)
1244
1245

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

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

1248
1249
class SafeWebBaseLoader(WebBaseLoader):
    """WebBaseLoader with enhanced error handling for URLs."""
Timothy J. Baek's avatar
Timothy J. Baek committed
1250

1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
    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
1263
1264
1265
                    metadata["description"] = description.get(
                        "content", "No description found."
                    )
1266
1267
                if html := soup.find("html"):
                    metadata["language"] = html.get("lang", "No language found.")
Timothy J. Baek's avatar
Timothy J. Baek committed
1268

1269
1270
1271
1272
                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
1273
1274


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
1275
1276
1277
1278
1279
1280
1281
1282
1283
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)}