main.py 14.8 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
4
5
6
7
8
9
from fastapi import (
    FastAPI,
    Depends,
    HTTPException,
    status,
    UploadFile,
    File,
    Form,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
10
from fastapi.middleware.cors import CORSMiddleware
Timothy J. Baek's avatar
Timothy J. Baek committed
11
import os, shutil
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 import embedding_functions
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,
Timothy J. Baek's avatar
Timothy J. Baek committed
23
    Docx2txtLoader,
Dave Bauman's avatar
Dave Bauman committed
24
    UnstructuredEPubLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
25
26
    UnstructuredWordDocumentLoader,
    UnstructuredMarkdownLoader,
27
    UnstructuredXMLLoader,
Marclass's avatar
Marclass committed
28
    UnstructuredRSTLoader,
Marclass's avatar
Marclass committed
29
    UnstructuredExcelLoader,
Timothy J. Baek's avatar
Timothy J. Baek committed
30
)
31
32
33
34
from langchain.text_splitter import RecursiveCharacterTextSplitter

from pydantic import BaseModel
from typing import Optional
35
import mimetypes
36
37
import uuid

38
39
40
41
42
from apps.web.models.documents import (
    Documents,
    DocumentForm,
    DocumentResponse,
)
Jannik Streidl's avatar
Jannik Streidl committed
43

44
45
46
47
48
49
from utils.misc import (
    calculate_sha256,
    calculate_sha256_string,
    sanitize_filename,
    extract_folders_after_data_docs,
)
50
from utils.utils import get_current_user, get_admin_user
51
52
53
from config import (
    UPLOAD_DIR,
    DOCS_DIR,
54
    SENTENCE_TRANSFORMER_EMBED_MODEL,
55
56
57
    CHROMA_CLIENT,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
Timothy J. Baek's avatar
Timothy J. Baek committed
58
    RAG_TEMPLATE,
59
)
60

61
62
from constants import ERROR_MESSAGES

63
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=SENTENCE_TRANSFORMER_EMBED_MODEL)
Timothy J. Baek's avatar
Timothy J. Baek committed
64
65
66

app = FastAPI()

Timothy J. Baek's avatar
Timothy J. Baek committed
67
68
app.state.CHUNK_SIZE = CHUNK_SIZE
app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
Timothy J. Baek's avatar
Timothy J. Baek committed
69
70
app.state.RAG_TEMPLATE = RAG_TEMPLATE

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

Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
76
77
78
79
80
81
82
origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
83
class CollectionNameForm(BaseModel):
84
85
86
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
89
90
class StoreWebForm(CollectionNameForm):
    url: str


Timothy J. Baek's avatar
Timothy J. Baek committed
91
def store_data_in_vector_db(data, collection_name) -> bool:
92
    text_splitter = RecursiveCharacterTextSplitter(
Timothy J. Baek's avatar
Timothy J. Baek committed
93
        chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP
94
95
96
97
98
99
    )
    docs = text_splitter.split_documents(data)

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

Timothy J. Baek's avatar
Timothy J. Baek committed
100
    try:
101
102
103
104
105
106
107
        if 'DOCKER_SENTENCE_TRANSFORMER_EMBED_MODEL' in os.environ:
    # if you use docker use the model from the environment variable
            collection = CHROMA_CLIENT.create_collection(name=collection_name, embedding_function=sentence_transformer_ef)
        
        else:
    # for local development use the default model
            collection = CHROMA_CLIENT.create_collection(name=collection_name)
108

Timothy J. Baek's avatar
Timothy J. Baek committed
109
        collection.add(
110
111
        documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts]
         )
Timothy J. Baek's avatar
Timothy J. Baek committed
112
113
114
115
116
117
118
        return True
    except Exception as e:
        print(e)
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False
119
120


Timothy J. Baek's avatar
Timothy J. Baek committed
121
122
@app.get("/")
async def get_status():
Timothy J. Baek's avatar
Timothy J. Baek committed
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
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
    }


@app.get("/chunk")
async def get_chunk_params(user=Depends(get_admin_user)):
    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
    }


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


@app.post("/chunk/update")
async def update_chunk_params(
    form_data: ChunkParamUpdateForm, user=Depends(get_admin_user)
):
    app.state.CHUNK_SIZE = form_data.chunk_size
    app.state.CHUNK_OVERLAP = form_data.chunk_overlap

    return {
        "status": True,
        "chunk_size": app.state.CHUNK_SIZE,
        "chunk_overlap": app.state.CHUNK_OVERLAP,
    }
156
157


Timothy J. Baek's avatar
Timothy J. Baek committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
@app.get("/template")
async def get_rag_template(user=Depends(get_current_user)):
    return {
        "status": True,
        "template": app.state.RAG_TEMPLATE,
    }


class RAGTemplateForm(BaseModel):
    template: str


@app.post("/template/update")
async def update_rag_template(form_data: RAGTemplateForm, user=Depends(get_admin_user)):
    # TODO: check template requirements
Timothy J. Baek's avatar
Timothy J. Baek committed
173
174
175
    app.state.RAG_TEMPLATE = (
        form_data.template if form_data.template != "" else RAG_TEMPLATE
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
176
    return {"status": True, "template": app.state.RAG_TEMPLATE}
177
178


179
class QueryDocForm(BaseModel):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
180
181
182
183
184
    collection_name: str
    query: str
    k: Optional[int] = 4


185
186
187
@app.post("/query/doc")
def query_doc(
    form_data: QueryDocForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
188
189
    user=Depends(get_current_user),
):
190
    try:
191
192
193
194
        if 'DOCKER_SENTENCE_TRANSFORMER_EMBED_MODEL' in os.environ:
        # if you use docker use the model from the environment variable
            collection = CHROMA_CLIENT.get_collection(
                name=form_data.collection_name,
Jannik Streidl's avatar
Jannik Streidl committed
195
                embedding_function=sentence_transformer_ef,
196
197
198
            )
        else:
        # for local development use the default model
Jannik Streidl's avatar
Jannik Streidl committed
199
200
            collection = CHROMA_CLIENT.get_collection(
            name=form_data.collection_name,
201
            )
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
202
        result = collection.query(query_texts=[form_data.query], n_results=form_data.k)
203
204
205
206
207
208
209
        return result
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
210
211


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
212
213
214
215
216
217
class QueryCollectionsForm(BaseModel):
    collection_names: List[str]
    query: str
    k: Optional[int] = 4


Timothy J. Baek's avatar
Timothy J. Baek committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def merge_and_sort_query_results(query_results, k):
    # Initialize lists to store combined data
    combined_ids = []
    combined_distances = []
    combined_metadatas = []
    combined_documents = []

    # Combine data from each dictionary
    for data in query_results:
        combined_ids.extend(data["ids"][0])
        combined_distances.extend(data["distances"][0])
        combined_metadatas.extend(data["metadatas"][0])
        combined_documents.extend(data["documents"][0])

    # Create a list of tuples (distance, id, metadata, document)
    combined = list(
        zip(combined_distances, combined_ids, combined_metadatas, combined_documents)
    )

    # Sort the list based on distances
    combined.sort(key=lambda x: x[0])

    # Unzip the sorted list
    sorted_distances, sorted_ids, sorted_metadatas, sorted_documents = zip(*combined)

    # Slicing the lists to include only k elements
    sorted_distances = list(sorted_distances)[:k]
    sorted_ids = list(sorted_ids)[:k]
    sorted_metadatas = list(sorted_metadatas)[:k]
    sorted_documents = list(sorted_documents)[:k]

    # Create the output dictionary
    merged_query_results = {
        "ids": [sorted_ids],
        "distances": [sorted_distances],
        "metadatas": [sorted_metadatas],
        "documents": [sorted_documents],
        "embeddings": None,
        "uris": None,
        "data": None,
    }

    return merged_query_results


263
264
@app.post("/query/collection")
def query_collection(
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
265
266
267
268
269
270
271
    form_data: QueryCollectionsForm,
    user=Depends(get_current_user),
):
    results = []

    for collection_name in form_data.collection_names:
        try:
272
            if 'DOCKER_SENTENCE_TRANSFORMER_EMBED_MODEL' in os.environ:
Jannik Streidl's avatar
Jannik Streidl committed
273
            # if you use docker use the model from the environment variable
274
                collection = CHROMA_CLIENT.get_collection(
Jannik Streidl's avatar
Jannik Streidl committed
275
276
                    name=collection_name,
                    embedding_function=sentence_transformer_ef,
277
278
279
280
                )
            else:
            # for local development use the default model
                collection = CHROMA_CLIENT.get_collection(
Jannik Streidl's avatar
Jannik Streidl committed
281
282
                name=collection_name,
                )
283
                
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
284
285
286
287
288
289
290
            result = collection.query(
                query_texts=[form_data.query], n_results=form_data.k
            )
            results.append(result)
        except:
            pass

Timothy J. Baek's avatar
Timothy J. Baek committed
291
    return merge_and_sort_query_results(results, form_data.k)
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
292
293


294
@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
295
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
296
297
298
299
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
        loader = WebBaseLoader(form_data.url)
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
300
301
302
303
304
305

        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)
Timothy J. Baek's avatar
Timothy J. Baek committed
306
307
        return {
            "status": True,
Timothy J. Baek's avatar
Timothy J. Baek committed
308
            "collection_name": collection_name,
Timothy J. Baek's avatar
Timothy J. Baek committed
309
310
            "filename": form_data.url,
        }
311
312
313
314
315
316
317
318
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


319
320
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
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
    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":
        loader = PyPDFLoader(file_path)
    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)
    elif file_ext == "md":
        loader = UnstructuredMarkdownLoader(file_path)
378
    elif file_content_type == "application/epub+zip":
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
379
380
        loader = UnstructuredEPubLoader(file_path)
    elif (
381
        file_content_type
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
382
383
384
385
        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        or file_ext in ["doc", "docx"]
    ):
        loader = Docx2txtLoader(file_path)
386
    elif file_content_type in [
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
387
388
389
390
        "application/vnd.ms-excel",
        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ] or file_ext in ["xls", "xlsx"]:
        loader = UnstructuredExcelLoader(file_path)
391
    elif file_ext in known_source_ext or file_content_type.find("text/") >= 0:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
392
393
394
395
396
397
398
399
        loader = TextLoader(file_path)
    else:
        loader = TextLoader(file_path)
        known_type = False

    return loader, known_type


400
@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
401
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
402
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
403
404
405
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
406
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
407

Timothy J. Baek's avatar
Timothy J. Baek committed
408
    print(file.content_type)
409
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
410
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
411
        file_path = f"{UPLOAD_DIR}/{filename}"
412
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
413
        with open(file_path, "wb") as f:
414
415
416
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
417
418
419
420
421
        f = open(file_path, "rb")
        if collection_name == None:
            collection_name = calculate_sha256(f)[:63]
        f.close()

422
        loader, known_type = get_loader(file.filename, file.content_type, file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
423
        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
424
425
426
        result = store_data_in_vector_db(data, collection_name)

        if result:
Timothy J. Baek's avatar
Timothy J. Baek committed
427
428
429
430
            return {
                "status": True,
                "collection_name": collection_name,
                "filename": filename,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
431
                "known_type": known_type,
Timothy J. Baek's avatar
Timothy J. Baek committed
432
            }
Timothy J. Baek's avatar
Timothy J. Baek committed
433
434
435
436
437
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
438
439
    except Exception as e:
        print(e)
Dave Bauman's avatar
Dave Bauman committed
440
441
442
443
444
445
446
447
448
449
        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),
            )
450
451


452
453
454
455
456
457
458
459
460
461
462
463
464
@app.get("/scan")
def scan_docs_dir(user=Depends(get_admin_user)):
    try:
        for path in Path(DOCS_DIR).rglob("./**/*"):
            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
465
466
467
                loader, known_type = get_loader(
                    filename, file_content_type[0], str(path)
                )
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
                data = loader.load()

                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,
                                                    )
                                                )
                                            }
                                        )
                                        if len(tags)
                                        else "{}"
                                    ),
                                }
                            ),
                        )

    except Exception as e:
        print(e)

    return True


Timothy J. Baek's avatar
Timothy J. Baek committed
509
@app.get("/reset/db")
510
511
def reset_vector_db(user=Depends(get_admin_user)):
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
512
513
514


@app.get("/reset")
515
516
517
518
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
519
        try:
520
521
522
523
            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
524
        except Exception as e:
525
            print("Failed to delete %s. Reason: %s" % (file_path, e))
Timothy J. Baek's avatar
Timothy J. Baek committed
526

527
528
529
530
531
532
    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
        print(e)

    return True