main.py 4.67 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

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
from langchain.document_loaders import WebBaseLoader, TextLoader, PyPDFLoader
17
18
19
20
21
22
23
24
25
26
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

Timothy J. Baek's avatar
Timothy J. Baek committed
27
from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
28
29
30
31
32
from constants import ERROR_MESSAGES

EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=EMBED_MODEL
)
Timothy J. Baek's avatar
Timothy J. Baek committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46

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
47
class CollectionNameForm(BaseModel):
48
49
50
    collection_name: Optional[str] = "test"


Timothy J. Baek's avatar
Timothy J. Baek committed
51
52
53
54
class StoreWebForm(CollectionNameForm):
    url: str


Timothy J. Baek's avatar
Timothy J. Baek committed
55
def store_data_in_vector_db(data, collection_name) -> bool:
56
57
58
59
60
61
62
63
    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
64
65
66
67
    try:
        collection = CHROMA_CLIENT.create_collection(
            name=collection_name, embedding_function=EMBEDDING_FUNC
        )
68

Timothy J. Baek's avatar
Timothy J. Baek committed
69
70
71
72
73
74
75
76
77
78
        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
79
80


Timothy J. Baek's avatar
Timothy J. Baek committed
81
82
83
@app.get("/")
async def get_status():
    return {"status": True}
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102


@app.get("/query/{collection_name}")
def query_collection(collection_name: str, query: str, k: Optional[int] = 4):
    collection = CHROMA_CLIENT.get_collection(
        name=collection_name,
    )
    result = collection.query(query_texts=[query], n_results=k)

    return result


@app.post("/web")
def store_web(form_data: StoreWebForm):
    # "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)
103
        return {"status": True, "collection_name": form_data.collection_name}
104
105
106
107
108
109
110
111
112
    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
113
def store_doc(collection_name: str = Form(...), file: UploadFile = File(...)):
114
    # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
115
    file.filename = f"{collection_name}-{file.filename}"
Timothy J. Baek's avatar
Timothy J. Baek committed
116
117
118
119
120
121
122

    if file.content_type not in ["application/pdf", "text/plain"]:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.FILE_NOT_SUPPORTED,
        )

123
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
124
        filename = file.filename
Timothy J. Baek's avatar
Timothy J. Baek committed
125
        file_path = f"{UPLOAD_DIR}/{filename}"
126
        contents = file.file.read()
Timothy J. Baek's avatar
Timothy J. Baek committed
127
        with open(file_path, "wb") as f:
128
129
130
            f.write(contents)
            f.close()

Timothy J. Baek's avatar
Timothy J. Baek committed
131
132
133
134
135
136
        if file.content_type == "application/pdf":
            loader = PyPDFLoader(file_path)
        elif file.content_type == "text/plain":
            loader = TextLoader(file_path)

        data = loader.load()
Timothy J. Baek's avatar
Timothy J. Baek committed
137
138
139
140
141
142
143
144
145
        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(),
            )
146
147
148
149
150
151
152
153
    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
154
@app.get("/reset/db")
155
156
def reset_vector_db():
    CHROMA_CLIENT.reset()
Timothy J. Baek's avatar
Timothy J. Baek committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176


@app.get("/reset")
def reset():
    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))

    try:
        CHROMA_CLIENT.reset()
    except Exception as e:
        print(e)

177
    return {"status": True}