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

Timothy J. Baek's avatar
Timothy J. Baek committed
14
# from chromadb.utils import embedding_functions
Timothy J. Baek's avatar
Timothy J. Baek committed
15

Timothy J. Baek's avatar
Timothy J. Baek committed
16
17
18
19
20
21
from langchain_community.document_loaders import (
    WebBaseLoader,
    TextLoader,
    PyPDFLoader,
    CSVLoader,
)
22
23
24
25
26
27
28
29
30
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA


from pydantic import BaseModel
from typing import Optional

import uuid
31
import time
32

Timothy J. Baek's avatar
Timothy J. Baek committed
33
from utils.misc import calculate_sha256
Timothy J. Baek's avatar
Timothy J. Baek committed
34
from utils.utils import get_current_user
Timothy J. Baek's avatar
Timothy J. Baek committed
35
from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
36
37
from constants import ERROR_MESSAGES

Timothy J. Baek's avatar
Timothy J. Baek committed
38
39
40
# EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
#     model_name=EMBED_MODEL
# )
Timothy J. Baek's avatar
Timothy J. Baek committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54

app = FastAPI()

origins = ["*"]

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


Timothy J. Baek's avatar
Timothy J. Baek committed
55
class CollectionNameForm(BaseModel):
56
57
58
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
59
60
61
62
class StoreWebForm(CollectionNameForm):
    url: str


Timothy J. Baek's avatar
Timothy J. Baek committed
63
def store_data_in_vector_db(data, collection_name) -> bool:
64
65
66
67
68
69
70
71
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
    )
    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
72
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
73
        collection = CHROMA_CLIENT.create_collection(name=collection_name)
74

Timothy J. Baek's avatar
Timothy J. Baek committed
75
76
77
78
79
80
81
82
83
84
        collection.add(
            documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts]
        )
        return True
    except Exception as e:
        print(e)
        if e.__class__.__name__ == "UniqueConstraintError":
            return True

        return False
85
86


Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
89
@app.get("/")
async def get_status():
    return {"status": True}
90
91
92


@app.get("/query/{collection_name}")
Timothy J. Baek's avatar
Timothy J. Baek committed
93
94
95
96
97
98
def query_collection(
    collection_name: str,
    query: str,
    k: Optional[int] = 4,
    user=Depends(get_current_user),
):
99
100
101
102
103
    try:
        collection = CHROMA_CLIENT.get_collection(
            name=collection_name,
        )
        result = collection.query(query_texts=[query], n_results=k)
104

105
106
107
108
109
110
111
        return result
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )
112
113
114


@app.post("/web")
Timothy J. Baek's avatar
Timothy J. Baek committed
115
def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
116
117
118
119
120
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
    try:
        loader = WebBaseLoader(form_data.url)
        data = loader.load()
        store_data_in_vector_db(data, form_data.collection_name)
121
        return {"status": True, "collection_name": form_data.collection_name}
122
123
124
125
126
127
128
129
130
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


@app.post("/doc")
Timothy J. Baek's avatar
Timothy J. Baek committed
131
def store_doc(
Timothy J. Baek's avatar
Timothy J. Baek committed
132
    collection_name: Optional[str] = Form(None),
Timothy J. Baek's avatar
Timothy J. Baek committed
133
134
135
    file: UploadFile = File(...),
    user=Depends(get_current_user),
):
136
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
Timothy J. Baek committed
137

Timothy J. Baek's avatar
Timothy J. Baek committed
138
    if file.content_type not in ["application/pdf", "text/plain", "text/csv"]:
Timothy J. Baek's avatar
Timothy J. Baek committed
139
140
141
142
143
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
        )

144
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
145
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
146
        file_path = f"{UPLOAD_DIR}/{filename}"
147
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
148
        with open(file_path, "wb") as f:
149
150
151
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
152
153
154
155
156
        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
157
158
159
160
        if file.content_type == "application/pdf":
            loader = PyPDFLoader(file_path)
        elif file.content_type == "text/plain":
            loader = TextLoader(file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
161
162
        elif file.content_type == "text/csv":
            loader = CSVLoader(file_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
163
164

        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
165
166
167
168
169
170
171
172
173
        result = store_data_in_vector_db(data, collection_name)

        if result:
            return {"status": True, "collection_name": collection_name}
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=ERROR_MESSAGES.DEFAULT(),
            )
174
175
176
177
178
179
180
181
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e),
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
182
@app.get("/reset/db")
Timothy J. Baek's avatar
Timothy J. Baek committed
183
184
185
186
187
188
189
190
def reset_vector_db(user=Depends(get_current_user)):
    if user.role == "admin":
        CHROMA_CLIENT.reset()
    else:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
191
192
193


@app.get("/reset")
194
def reset(user=Depends(get_current_user)) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
195
196
197
198
199
200
201
202
203
204
205
206
    if user.role == "admin":
        folder = f"{UPLOAD_DIR}"
        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)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print("Failed to delete %s. Reason: %s" % (file_path, e))

Timothy J. Baek's avatar
Timothy J. Baek committed
207
        try:
Timothy J. Baek's avatar
Timothy J. Baek committed
208
            CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
209
        except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
210
            print(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
211

212
        return True
Timothy J. Baek's avatar
Timothy J. Baek committed
213
214
215
216
217
    else:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )