refac: file upload

This commit is contained in:
Timothy J. Baek 2024-01-07 09:00:30 -08:00
parent ffd0a5a2a0
commit d6a1bf1406
2 changed files with 11 additions and 7 deletions

View file

@ -24,7 +24,7 @@ from typing import Optional
import uuid import uuid
from utils.misc import calculate_sha256
from utils.utils import get_current_user from utils.utils import get_current_user
from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP from config import UPLOAD_DIR, EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
from constants import ERROR_MESSAGES from constants import ERROR_MESSAGES
@ -123,12 +123,11 @@ def store_web(form_data: StoreWebForm, user=Depends(get_current_user)):
@app.post("/doc") @app.post("/doc")
def store_doc( def store_doc(
collection_name: str = Form(...), collection_name: Optional[str] = Form(None),
file: UploadFile = File(...), file: UploadFile = File(...),
user=Depends(get_current_user), user=Depends(get_current_user),
): ):
# "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"{collection_name}-{file.filename}"
if file.content_type not in ["application/pdf", "text/plain"]: if file.content_type not in ["application/pdf", "text/plain"]:
raise HTTPException( raise HTTPException(
@ -144,6 +143,11 @@ def store_doc(
f.write(contents) f.write(contents)
f.close() f.close()
f = open(file_path, "rb")
if collection_name == None:
collection_name = calculate_sha256(f)[:63]
f.close()
if file.content_type == "application/pdf": if file.content_type == "application/pdf":
loader = PyPDFLoader(file_path) loader = PyPDFLoader(file_path)
elif file.content_type == "text/plain": elif file.content_type == "text/plain":

View file

@ -124,8 +124,8 @@
reader.readAsDataURL(file); reader.readAsDataURL(file);
} else if (['application/pdf', 'text/plain'].includes(file['type'])) { } else if (['application/pdf', 'text/plain'].includes(file['type'])) {
console.log(file); console.log(file);
const hash = (await calculateSHA256(file)).substring(0, 63); // const hash = (await calculateSHA256(file)).substring(0, 63);
const res = await uploadDocToVectorDB(localStorage.token, hash, file); const res = await uploadDocToVectorDB(localStorage.token, '', file);
if (res) { if (res) {
files = [ files = [
@ -243,8 +243,8 @@
reader.readAsDataURL(file); reader.readAsDataURL(file);
} else if (['application/pdf', 'text/plain'].includes(file['type'])) { } else if (['application/pdf', 'text/plain'].includes(file['type'])) {
console.log(file); console.log(file);
const hash = (await calculateSHA256(file)).substring(0, 63); // const hash = (await calculateSHA256(file)).substring(0, 63);
const res = await uploadDocToVectorDB(localStorage.token, hash, file); const res = await uploadDocToVectorDB(localStorage.token, '', file);
if (res) { if (res) {
files = [ files = [