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

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

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

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

35
36
37
38
39
import validators
import urllib.parse
import socket


40
41
from pydantic import BaseModel
from typing import Optional
42
import mimetypes
43
import uuid
44
45
import json

46
import sentence_transformers
47

48
49
50
51
52
from apps.web.models.documents import (
    Documents,
    DocumentForm,
    DocumentResponse,
)
Jannik Streidl's avatar
Jannik Streidl committed
53

54
from apps.rag.utils import (
55
    get_model_path,
Timothy J. Baek's avatar
Timothy J. Baek committed
56
57
58
59
60
    get_embedding_function,
    query_doc,
    query_doc_with_hybrid_search,
    query_collection,
    query_collection_with_hybrid_search,
61
)
Timothy J. Baek's avatar
Timothy J. Baek committed
62

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

71
from config import (
72
    SRC_LOG_LEVELS,
73
74
    UPLOAD_DIR,
    DOCS_DIR,
75
76
    RAG_TOP_K,
    RAG_RELEVANCE_THRESHOLD,
77
    RAG_EMBEDDING_ENGINE,
78
    RAG_EMBEDDING_MODEL,
79
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
80
    RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
81
    ENABLE_RAG_HYBRID_SEARCH,
Steven Kreitzer's avatar
Steven Kreitzer committed
82
    RAG_RERANKING_MODEL,
83
    PDF_EXTRACT_IMAGES,
84
    RAG_RERANKING_MODEL_AUTO_UPDATE,
Steven Kreitzer's avatar
Steven Kreitzer committed
85
    RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
Timothy J. Baek's avatar
Timothy J. Baek committed
86
87
    RAG_OPENAI_API_BASE_URL,
    RAG_OPENAI_API_KEY,
88
    DEVICE_TYPE,
89
90
91
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
92
    RAG_TEMPLATE,
93
    ENABLE_LOCAL_WEB_FETCH,
94
)
95

96
97
from constants import ERROR_MESSAGES

98
99
100
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])

Timothy J. Baek's avatar
Timothy J. Baek committed
101
102
app = FastAPI()

103
104
app.state.TOP_K = RAG_TOP_K
app.state.RELEVANCE_THRESHOLD = RAG_RELEVANCE_THRESHOLD
Timothy J. Baek's avatar
Timothy J. Baek committed
105
106

app.state.ENABLE_RAG_HYBRID_SEARCH = ENABLE_RAG_HYBRID_SEARCH
Steven Kreitzer's avatar
Steven Kreitzer committed
107

Timothy J. Baek's avatar
Timothy J. Baek committed
108
109
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
110

111
app.state.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
112
app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
Steven Kreitzer's avatar
Steven Kreitzer committed
113
app.state.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
114
app.state.RAG_TEMPLATE = RAG_TEMPLATE
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
115

Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
116
117
app.state.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = RAG_OPENAI_API_KEY
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
118

119
app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
120

Steven Kreitzer's avatar
Steven Kreitzer committed
121

122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def update_embedding_model(
    embedding_model: str,
    update_model: bool = False,
):
    if embedding_model and app.state.RAG_EMBEDDING_ENGINE == "":
        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(
    app.state.RAG_EMBEDDING_MODEL,
    RAG_EMBEDDING_MODEL_AUTO_UPDATE,
)

update_reranking_model(
    app.state.RAG_RERANKING_MODEL,
    RAG_RERANKING_MODEL_AUTO_UPDATE,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
159

Timothy J. Baek's avatar
Timothy J. Baek committed
160
161
162
163
164
165
166
167
168

app.state.EMBEDDING_FUNCTION = get_embedding_function(
    app.state.RAG_EMBEDDING_ENGINE,
    app.state.RAG_EMBEDDING_MODEL,
    app.state.sentence_transformer_ef,
    app.state.OPENAI_API_KEY,
    app.state.OPENAI_API_BASE_URL,
)

Timothy J. Baek's avatar
Timothy J. Baek committed
169
170
origins = ["*"]

171

Timothy J. Baek's avatar
Timothy J. Baek committed
172
173
174
175
176
177
178
179
180
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
181
class CollectionNameForm(BaseModel):
182
183
184
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
185
class UrlForm(CollectionNameForm):
Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
    url: str

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

Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
191
192
193
194
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
195
        "template": app.state.RAG_TEMPLATE,
196
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
197
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
Steven Kreitzer's avatar
Steven Kreitzer committed
198
        "reranking_model": app.state.RAG_RERANKING_MODEL,
199
200
201
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
202
203
@app.get("/embedding")
async def get_embedding_config(user=Depends(get_admin_user)):
204
205
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
206
        "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
207
        "embedding_model": app.state.RAG_EMBEDDING_MODEL,
208
        "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
209
210
            "url": app.state.OPENAI_API_BASE_URL,
            "key": app.state.OPENAI_API_KEY,
211
        },
212
213
214
    }


Steven Kreitzer's avatar
Steven Kreitzer committed
215
216
217
218
219
@app.get("/reranking")
async def get_reraanking_config(user=Depends(get_admin_user)):
    return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}


220
221
222
223
224
class OpenAIConfigForm(BaseModel):
    url: str
    key: str


225
class EmbeddingModelUpdateForm(BaseModel):
226
    openai_config: Optional[OpenAIConfigForm] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
227
    embedding_engine: str
228
229
230
    embedding_model: str


Timothy J. Baek's avatar
Timothy J. Baek committed
231
232
@app.post("/embedding/update")
async def update_embedding_config(
233
234
    form_data: EmbeddingModelUpdateForm, user=Depends(get_admin_user)
):
Self Denial's avatar
Self Denial committed
235
236
    log.info(
        f"Updating embedding model: {app.state.RAG_EMBEDDING_MODEL} to {form_data.embedding_model}"
237
    )
238
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
239
        app.state.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
240
        app.state.RAG_EMBEDDING_MODEL = form_data.embedding_model
Timothy J. Baek's avatar
Timothy J. Baek committed
241

242
243
        if app.state.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
            if form_data.openai_config != None:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
244
245
                app.state.OPENAI_API_BASE_URL = form_data.openai_config.url
                app.state.OPENAI_API_KEY = form_data.openai_config.key
246

247
        update_embedding_model(app.state.RAG_EMBEDDING_MODEL, True)
248

Timothy J. Baek's avatar
Timothy J. Baek committed
249
250
251
252
253
254
255
256
        app.state.EMBEDDING_FUNCTION = get_embedding_function(
            app.state.RAG_EMBEDDING_ENGINE,
            app.state.RAG_EMBEDDING_MODEL,
            app.state.sentence_transformer_ef,
            app.state.OPENAI_API_KEY,
            app.state.OPENAI_API_BASE_URL,
        )

257
258
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
259
            "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
260
            "embedding_model": app.state.RAG_EMBEDDING_MODEL,
261
            "openai_config": {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
262
263
                "url": app.state.OPENAI_API_BASE_URL,
                "key": app.state.OPENAI_API_KEY,
264
            },
265
266
267
268
269
270
271
        }
    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
272
273


Steven Kreitzer's avatar
Steven Kreitzer committed
274
275
class RerankingModelUpdateForm(BaseModel):
    reranking_model: str
276

Steven Kreitzer's avatar
Steven Kreitzer committed
277
278
279
280
281
282
283
284
285
286

@app.post("/reranking/update")
async def update_reranking_config(
    form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
):
    log.info(
        f"Updating reranking model: {app.state.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
    )
    try:
        app.state.RAG_RERANKING_MODEL = form_data.reranking_model
287

288
        update_reranking_model(app.state.RAG_RERANKING_MODEL, True)
Steven Kreitzer's avatar
Steven Kreitzer committed
289
290
291
292
293
294
295
296
297
298
299
300
301

        return {
            "status": True,
            "reranking_model": app.state.RAG_RERANKING_MODEL,
        }
    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
302
303
@app.get("/config")
async def get_rag_config(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
304
305
    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
306
307
308
309
310
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
311
312
313
314
315
316
317
318
    }


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


Timothy J. Baek's avatar
Timothy J. Baek committed
319
320
321
322
323
324
325
326
327
328
class ConfigUpdateForm(BaseModel):
    pdf_extract_images: bool
    chunk: ChunkParamUpdateForm


@app.post("/config/update")
async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
    app.state.PDF_EXTRACT_IMAGES = form_data.pdf_extract_images
    app.state.CHUNK_SIZE = form_data.chunk.chunk_size
    app.state.CHUNK_OVERLAP = form_data.chunk.chunk_overlap
Timothy J. Baek's avatar
Timothy J. Baek committed
329
330
331

    return {
        "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
332
333
334
335
336
        "pdf_extract_images": app.state.PDF_EXTRACT_IMAGES,
        "chunk": {
            "chunk_size": app.state.CHUNK_SIZE,
            "chunk_overlap": app.state.CHUNK_OVERLAP,
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
337
    }
338
339


Timothy J. Baek's avatar
Timothy J. Baek committed
340
341
342
343
344
345
346
347
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


348
349
350
351
352
353
@app.get("/query/settings")
async def get_query_settings(user=Depends(get_admin_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
        "k": app.state.TOP_K,
354
        "r": app.state.RELEVANCE_THRESHOLD,
Timothy J. Baek's avatar
Timothy J. Baek committed
355
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
356
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
357
358


359
360
class QuerySettingsForm(BaseModel):
    k: Optional[int] = None
361
    r: Optional[float] = None
362
    template: Optional[str] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
363
    hybrid: Optional[bool] = None
364
365
366
367
368
369
370
371


@app.post("/query/settings/update")
async def update_query_settings(
    form_data: QuerySettingsForm, user=Depends(get_admin_user)
):
    app.state.RAG_TEMPLATE = form_data.template if form_data.template else RAG_TEMPLATE
    app.state.TOP_K = form_data.k if form_data.k else 4
372
    app.state.RELEVANCE_THRESHOLD = form_data.r if form_data.r else 0.0
Timothy J. Baek's avatar
Timothy J. Baek committed
373
    app.state.ENABLE_RAG_HYBRID_SEARCH = form_data.hybrid if form_data.hybrid else False
Steven Kreitzer's avatar
Steven Kreitzer committed
374
375
376
377
378
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
        "k": app.state.TOP_K,
        "r": app.state.RELEVANCE_THRESHOLD,
Timothy J. Baek's avatar
Timothy J. Baek committed
379
        "hybrid": app.state.ENABLE_RAG_HYBRID_SEARCH,
Steven Kreitzer's avatar
Steven Kreitzer committed
380
    }
381
382


383
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
384
385
    collection_name: str
    query: str
386
    k: Optional[int] = None
387
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
388
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
389
390


391
@app.post("/query/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
392
def query_doc_handler(
393
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
394
395
    user=Depends(get_current_user),
):
396
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
397
398
399
400
        if app.state.ENABLE_RAG_HYBRID_SEARCH:
            return query_doc_with_hybrid_search(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
401
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
402
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
403
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
404
405
406
407
408
409
                r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
            )
        else:
            return query_doc(
                collection_name=form_data.collection_name,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
410
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
411
412
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
413
    except Exception as e:
414
        log.exception(e)
415
416
417
418
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
419
420


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
421
422
423
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
424
    k: Optional[int] = None
425
    r: Optional[float] = None
Steven Kreitzer's avatar
Steven Kreitzer committed
426
    hybrid: Optional[bool] = None
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
427
428


429
@app.post("/query/collection")
Timothy J. Baek's avatar
Timothy J. Baek committed
430
def query_collection_handler(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
431
432
433
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
434
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
435
436
437
438
        if app.state.ENABLE_RAG_HYBRID_SEARCH:
            return query_collection_with_hybrid_search(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
439
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
440
                k=form_data.k if form_data.k else app.state.TOP_K,
Steven Kreitzer's avatar
Steven Kreitzer committed
441
                reranking_function=app.state.sentence_transformer_rf,
Timothy J. Baek's avatar
Timothy J. Baek committed
442
443
444
445
446
447
                r=form_data.r if form_data.r else app.state.RELEVANCE_THRESHOLD,
            )
        else:
            return query_collection(
                collection_names=form_data.collection_names,
                query=form_data.query,
Steven Kreitzer's avatar
Steven Kreitzer committed
448
                embedding_function=app.state.EMBEDDING_FUNCTION,
Timothy J. Baek's avatar
Timothy J. Baek committed
449
450
                k=form_data.k if form_data.k else app.state.TOP_K,
            )
451

452
453
454
455
456
457
    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
458
459


Timothy J. Baek's avatar
Timothy J. Baek committed
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
@app.post("/youtube")
def store_youtube_video(form_data: UrlForm, user=Depends(get_current_user)):
    try:
        loader = YoutubeLoader.from_youtube_url(form_data.url, add_video_info=False)
        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),
        )


484
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
485
def store_web(form_data: UrlForm, user=Depends(get_current_user)):
486
487
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
488
        loader = get_web_loader(form_data.url)
489
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
490
491
492
493
494

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

495
        store_data_in_vector_db(data, collection_name, overwrite=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
496
497
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
498
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
499
500
            "filename": form_data.url,
        }
501
    except Exception as e:
502
        log.exception(e)
503
504
505
506
507
508
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
def get_web_loader(url: str):
    # Check if the URL is valid
    if isinstance(validators.url(url), validators.ValidationError):
        raise ValueError(ERROR_MESSAGES.INVALID_URL)
    if not ENABLE_LOCAL_WEB_FETCH:
        # Local web fetch is disabled, filter out any URLs that resolve to private IP addresses
        parsed_url = urllib.parse.urlparse(url)
        # Get IPv4 and IPv6 addresses
        ipv4_addresses, ipv6_addresses = resolve_hostname(parsed_url.hostname)
        # Check if any of the resolved addresses are private
        # This is technically still vulnerable to DNS rebinding attacks, as we don't control WebBaseLoader
        for ip in ipv4_addresses:
            if validators.ipv4(ip, private=True):
                raise ValueError(ERROR_MESSAGES.INVALID_URL)
        for ip in ipv6_addresses:
            if validators.ipv6(ip, private=True):
                raise ValueError(ERROR_MESSAGES.INVALID_URL)
    return WebBaseLoader(url)


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


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

542
543
544
545
546
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
547

548
    docs = text_splitter.split_documents(data)
Timothy J. Baek's avatar
Timothy J. Baek committed
549
550

    if len(docs) > 0:
551
        log.info(f"store_data_in_vector_db {docs}")
Timothy J. Baek's avatar
Timothy J. Baek committed
552
553
554
        return store_docs_in_vector_db(docs, collection_name, overwrite), None
    else:
        raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
555
556
557


def store_text_in_vector_db(
Timothy J. Baek's avatar
Timothy J. Baek committed
558
    text, metadata, collection_name, overwrite: bool = False
559
560
561
562
563
564
) -> bool:
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=app.state.CHUNK_SIZE,
        chunk_overlap=app.state.CHUNK_OVERLAP,
        add_start_index=True,
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
565
    docs = text_splitter.create_documents([text], metadatas=[metadata])
566
567
568
    return store_docs_in_vector_db(docs, collection_name, overwrite)


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

572
573
574
575
576
577
578
    texts = [doc.page_content for doc in docs]
    metadatas = [doc.metadata for doc in docs]

    try:
        if overwrite:
            for collection in CHROMA_CLIENT.list_collections():
                if collection_name == collection.name:
579
                    log.info(f"deleting existing collection {collection_name}")
580
581
                    CHROMA_CLIENT.delete_collection(name=collection_name)

582
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
583

Timothy J. Baek's avatar
Timothy J. Baek committed
584
        embedding_func = get_embedding_function(
Steven Kreitzer's avatar
Steven Kreitzer committed
585
586
587
588
589
590
591
592
            app.state.RAG_EMBEDDING_ENGINE,
            app.state.RAG_EMBEDDING_MODEL,
            app.state.sentence_transformer_ef,
            app.state.OPENAI_API_KEY,
            app.state.OPENAI_API_BASE_URL,
        )

        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
593
        embeddings = embedding_func(embedding_texts)
594
595
596
597
598
599
600
601
602

        for batch in create_batches(
            api=CHROMA_CLIENT,
            ids=[str(uuid.uuid1()) for _ in texts],
            metadatas=metadatas,
            embeddings=embeddings,
            documents=texts,
        ):
            collection.add(*batch)
603

604
        return True
605
    except Exception as e:
606
        log.exception(e)
607
608
609
610
611
612
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False


613
614
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
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
    known_type = True

    known_source_ext = [
        "go",
        "py",
        "java",
        "sh",
        "bat",
        "ps1",
        "cmd",
        "js",
        "ts",
        "css",
        "cpp",
        "hpp",
        "h",
        "c",
        "cs",
        "sql",
        "log",
        "ini",
        "pl",
        "pm",
        "r",
        "dart",
        "dockerfile",
        "env",
        "php",
        "hs",
        "hsc",
        "lua",
        "nginxconf",
        "conf",
        "m",
        "mm",
        "plsql",
        "perl",
        "rb",
        "rs",
        "db2",
        "scala",
        "bash",
        "swift",
        "vue",
        "svelte",
    ]

    if file_ext == "pdf":
Timothy J. Baek's avatar
Timothy J. Baek committed
663
        loader = PyPDFLoader(file_path, extract_images=app.state.PDF_EXTRACT_IMAGES)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
664
665
666
667
668
669
    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)
670
    elif file_ext in ["htm", "html"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
671
        loader = BSHTMLLoader(file_path, open_encoding="unicode_escape")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
672
673
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
674
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
675
676
        loader = UnstructuredEPubLoader(file_path)
    elif (
677
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
678
679
680
681
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
682
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
683
684
685
686
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
687
688
689
    elif file_ext in known_source_ext or (
        file_content_type and file_content_type.find("text/") >= 0
    ):
690
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
691
    else:
692
        loader = TextLoader(file_path, autodetect_encoding=True)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
693
694
695
696
697
        known_type = False

    return loader, known_type


698
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
699
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
700
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
701
702
703
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
704
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
705

706
    log.info(f"file.content_type: {file.content_type}")
707
    try:
708
        unsanitized_filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
709
        filename = os.path.basename(unsanitized_filename)
710

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

713
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
714
        with open(file_path, "wb") as f:
715
716
717
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
718
719
720
721
722
        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
723
        loader, known_type = get_loader(filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
724
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
725
726
727
728
729
730
731
732
733
734
735
736

        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
737
738
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
Timothy J. Baek's avatar
Timothy J. Baek committed
739
                detail=e,
Timothy J. Baek's avatar
Timothy J. Baek committed
740
            )
741
    except Exception as e:
742
        log.exception(e)
Dave Bauman's avatar
Dave Bauman committed
743
744
745
746
747
748
749
750
751
752
        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),
            )
753
754


755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
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
771
772
773
774
775
    result = store_text_in_vector_db(
        form_data.content,
        metadata={"name": form_data.name, "created_by": user.id},
        collection_name=collection_name,
    )
776
777
778
779
780
781
782
783
784
785

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


786
787
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
788
789
    for path in Path(DOCS_DIR).rglob("./**/*"):
        try:
790
791
792
793
794
795
796
797
798
            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
799
800
801
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
802
803
                data = loader.load()

Timothy J. Baek's avatar
Timothy J. Baek committed
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
                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,
                                                        )
828
                                                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
829
830
831
832
833
834
835
836
837
                                                }
                                            )
                                            if len(tags)
                                            else "{}"
                                        ),
                                    }
                                ),
                            )
                except Exception as e:
838
                    log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
839
                    pass
840

841
        except Exception as e:
842
            log.exception(e)
843
844
845
846

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
847
@app.get("/reset/db")
848
849
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
850
851
852


@app.get("/reset")
853
854
855
856
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
857
        try:
858
859
860
861
            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
862
        except Exception as e:
863
            log.error("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
864

865
866
867
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
868
        log.exception(e)
869
870

    return True