Daniel Marques commited on
Commit
177debf
1 Parent(s): b1f4ef7

feat: add upload file

Browse files
Files changed (1) hide show
  1. main.py +39 -21
main.py CHANGED
@@ -149,7 +149,6 @@ def delete_source_route():
149
 
150
  return jsonify({"message": f"Folder '{folder_name}' successfully deleted and recreated."})
151
 
152
-
153
  @api_app.post('/predict')
154
  async def predict(data: Predict):
155
  global QA
@@ -176,26 +175,45 @@ async def predict(data: Predict):
176
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
177
 
178
  @api_app.post("/save_document/")
179
- async def create_upload_file(file: Union[UploadFile, None] = None):
180
- try:
181
- if not file:
182
- raise HTTPException(status_code=400, detail="No upload file sent")
183
- else:
184
- if file.filename == "":
185
- raise HTTPException(status_code=400, detail="No selected file")
186
- if file:
187
- filename = file.filename
188
- folder_path = "SOURCE_DOCUMENTS"
189
-
190
- if not os.path.exists(folder_path):
191
- os.makedirs(folder_path)
192
-
193
- file_path = os.path.join(folder_path, filename)
194
- file.save(file_path)
195
-
196
- return {"response": "File saved successfully"}
197
- except Exception as e:
198
- raise HTTPException(status_code=400, detail=e)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
199
 
200
  @api_app.websocket("/ws")
201
  async def websocket_endpoint(websocket: WebSocket):
 
149
 
150
  return jsonify({"message": f"Folder '{folder_name}' successfully deleted and recreated."})
151
 
 
152
  @api_app.post('/predict')
153
  async def predict(data: Predict):
154
  global QA
 
175
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
176
 
177
  @api_app.post("/save_document/")
178
+ async def create_upload_file(file: UploadFile):
179
+ # Get the file size (in bytes)
180
+ file.file.seek(0, 2)
181
+ file_size = file.file.tell()
182
+
183
+ # move the cursor back to the beginning
184
+ await file.seek(0)
185
+
186
+ if file_size > 2 * 1024 * 1024:
187
+ # more than 2 MB
188
+ raise HTTPException(status_code=400, detail="File too large")
189
+
190
+ # check the content type (MIME type)
191
+ content_type = file.content_type
192
+ if content_type not in ["image/jpeg", "image/png", "image/gif"]:
193
+ raise HTTPException(status_code=400, detail="Invalid file type")
194
+
195
+ # do something with the valid file
196
+ return {"filename": file.filename}
197
+ # async def create_upload_file(file: Union[UploadFile, None] = None):
198
+ # try:
199
+ # if not file:
200
+ # raise HTTPException(status_code=400, detail="No upload file sent")
201
+ # else:
202
+ # if file.filename == "":
203
+ # raise HTTPException(status_code=400, detail="No selected file")
204
+ # if file:
205
+ # filename = file.filename
206
+ # folder_path = "SOURCE_DOCUMENTS"
207
+
208
+ # if not os.path.exists(folder_path):
209
+ # os.makedirs(folder_path)
210
+
211
+ # file_path = os.path.join(folder_path, filename)
212
+ # file.save(file_path)
213
+
214
+ # return {"response": "File saved successfully"}
215
+ # except Exception as e:
216
+ # raise HTTPException(status_code=400, detail=e)
217
 
218
  @api_app.websocket("/ws")
219
  async def websocket_endpoint(websocket: WebSocket):