feat: pdf data load

This commit is contained in:
Timothy J. Baek 2024-01-06 23:40:51 -08:00
parent 784b369cc9
commit cd86c36953
3 changed files with 40 additions and 11 deletions

View file

@ -1,9 +1,18 @@
from fastapi import FastAPI, Request, Depends, HTTPException, status, UploadFile, File from fastapi import (
FastAPI,
Request,
Depends,
HTTPException,
status,
UploadFile,
File,
Form,
)
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from chromadb.utils import embedding_functions from chromadb.utils import embedding_functions
from langchain.document_loaders import WebBaseLoader, TextLoader from langchain.document_loaders import WebBaseLoader, TextLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA from langchain.chains import RetrievalQA
@ -34,11 +43,14 @@ app.add_middleware(
) )
class StoreWebForm(BaseModel): class CollectionNameForm(BaseModel):
url: str
collection_name: Optional[str] = "test" collection_name: Optional[str] = "test"
class StoreWebForm(CollectionNameForm):
url: str
def store_data_in_vector_db(data, collection_name): def store_data_in_vector_db(data, collection_name):
text_splitter = RecursiveCharacterTextSplitter( text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
@ -89,20 +101,34 @@ def store_web(form_data: StoreWebForm):
@app.post("/doc") @app.post("/doc")
def store_doc(file: UploadFile = File(...)): def store_doc(collection_name: str = Form(...), file: UploadFile = File(...)):
# "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm" # "https://www.gutenberg.org/files/1727/1727-h/1727-h.htm"
file.filename = f"{uuid.uuid4()}-{file.filename}"
print(dir(file))
print(file.content_type)
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,
)
try: try:
print(file) filename = file.filename
file.filename = f"{uuid.uuid4()}-{file.filename}" file_path = f"./data/{filename}"
contents = file.file.read() contents = file.file.read()
with open(f"./data/{file.filename}", "wb") as f: with open(file_path, "wb") as f:
f.write(contents) f.write(contents)
f.close() f.close()
# loader = WebBaseLoader(form_data.url) if file.content_type == "application/pdf":
# data = loader.load() loader = PyPDFLoader(file_path)
# store_data_in_vector_db(data, form_data.collection_name) elif file.content_type == "text/plain":
loader = TextLoader(file_path)
data = loader.load()
store_data_in_vector_db(data, collection_name)
return {"status": True} return {"status": True}
except Exception as e: except Exception as e:
print(e) print(e)

View file

@ -33,6 +33,8 @@ class ERROR_MESSAGES(str, Enum):
) )
FILE_NOT_SENT = "FILE_NOT_SENT" FILE_NOT_SENT = "FILE_NOT_SENT"
FILE_NOT_SUPPORTED = "Oops! It seems like the file format you're trying to upload is not supported. Please upload a file with a supported format (e.g., JPG, PNG, PDF, TXT) and try again."
NOT_FOUND = "We could not find what you're looking for :/" NOT_FOUND = "We could not find what you're looking for :/"
USER_NOT_FOUND = "We could not find what you're looking for :/" USER_NOT_FOUND = "We could not find what you're looking for :/"
API_KEY_NOT_FOUND = "Oops! It looks like there's a hiccup. The API key is missing. Please make sure to provide a valid API key to access this feature." API_KEY_NOT_FOUND = "Oops! It looks like there's a hiccup. The API key is missing. Please make sure to provide a valid API key to access this feature."

View file

@ -20,6 +20,7 @@ bcrypt
langchain langchain
chromadb chromadb
sentence_transformers sentence_transformers
pypdf
PyJWT PyJWT
pyjwt[crypto] pyjwt[crypto]