Daniel Marques commited on
Commit
d0c75d3
1 Parent(s): 19d7335

feat: add save document route

Browse files
Files changed (1) hide show
  1. main.py +21 -2
main.py CHANGED
@@ -1,4 +1,5 @@
1
- from fastapi import FastAPI, HTTPException
 
2
  from fastapi.staticfiles import StaticFiles
3
 
4
  from pydantic import BaseModel
@@ -50,7 +51,7 @@ RETRIEVER = DB.as_retriever()
50
  LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
51
  prompt, memory = get_prompt_template(promptTemplate_type="llama", history=False)
52
 
53
- template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer.
54
  {context}
55
  Question: {question}
56
  Helpful Answer:"""
@@ -146,3 +147,21 @@ def run_ingest_route():
146
  except Exception as e:
147
  raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
148
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List
2
+ from fastapi import FastAPI, HTTPException, UploadFile
3
  from fastapi.staticfiles import StaticFiles
4
 
5
  from pydantic import BaseModel
 
51
  LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
52
  prompt, memory = get_prompt_template(promptTemplate_type="llama", history=False)
53
 
54
+ template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use 10 sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer.
55
  {context}
56
  Question: {question}
57
  Helpful Answer:"""
 
147
  except Exception as e:
148
  raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
149
 
150
+ @app.post("/save_document/")
151
+ async def create_upload_file(file: Union[UploadFile, None] = None):
152
+ if not file:
153
+ raise HTTPException(status_code=400, detail="No upload file sent")
154
+ else:
155
+ if file.filename == "":
156
+ raise HTTPException(status_code=400, detail="No selected file")
157
+ if file:
158
+ filename = secure_filename(file.filename)
159
+ folder_path = "SOURCE_DOCUMENTS"
160
+
161
+ if not os.path.exists(folder_path):
162
+ os.makedirs(folder_path)
163
+
164
+ file_path = os.path.join(folder_path, filename)
165
+ file.save(file_path)
166
+
167
+ return {"response": "File saved successfully"}