feat: add save document route

appv1.py

@@ -26,8 +26,6 @@ else:
     DEVICE_TYPE = "cpu"
 
 SHOW_SOURCES = True
-logging.info(f"Running on: {DEVICE_TYPE}")
-logging.info(f"Display Source Documents set to: {SHOW_SOURCES}")
 
 EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": DEVICE_TYPE})
 
@@ -182,4 +180,4 @@ def prompt_route():
 
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5110)
+    app.run(host="0.0.0.0", port=5110)

main.py

@@ -1,7 +1,8 @@
 from typing import Union
-from fastapi import FastAPI, HTTPException, UploadFile
+from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
 from fastapi.staticfiles import StaticFiles
 
+
 from pydantic import BaseModel
 import pickle
 import uvicorn
@@ -26,16 +27,16 @@ from werkzeug.utils import secure_filename
 
 from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME
 
-if torch.backends.mps.is_available():
-    DEVICE_TYPE = "mps"
-elif torch.cuda.is_available():
-    DEVICE_TYPE = "cuda"
-else:
-    DEVICE_TYPE = "cpu"
+# if torch.backends.mps.is_available():
+#     DEVICE_TYPE = "mps"
+# elif torch.cuda.is_available():
+#     DEVICE_TYPE = "cuda"
+# else:
+#     DEVICE_TYPE = "cpu"
+
+DEVICE_TYPE = "cuda"
 
 SHOW_SOURCES = True
-logging.info(f"Running on: {DEVICE_TYPE}")
-logging.info(f"Display Source Documents set to: {SHOW_SOURCES}")
 
 EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": DEVICE_TYPE})
 
@@ -51,7 +52,14 @@ RETRIEVER = DB.as_retriever()
 LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
 prompt, memory = get_prompt_template(promptTemplate_type="llama", history=False)
 
-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.
+template = """you are a helpful, respectful and honest assistant.
+Your name is Katara llma. You should only use the source documents provided to answer the questions.
+You should only respond only topics that contains in documents use to training.
+Use the following pieces of context to answer the question at the end.
+Always answer in the most helpful and safe way possible.
+If you don't know the answer to a question, just say that you don't know, don't try to make up an answer, don't share false information.
+Use 15 sentences maximum. Keep the answer as concise as possible.
+Always say "thanks for asking!" at the end of the answer.
 {context}
 Question: {question}
 Helpful Answer:"""
@@ -149,19 +157,29 @@ def run_ingest_route():
 
 @api_app.post("/save_document/")
 async def create_upload_file(file: Union[UploadFile, None] = None):
-    if not file:
-        raise HTTPException(status_code=400, detail="No upload file sent")
-    else:
-        if file.filename == "":
-             raise HTTPException(status_code=400, detail="No selected file")
-        if file:
-            filename = secure_filename(file.filename)
-            folder_path = "SOURCE_DOCUMENTS"
+    try:
+        if not file:
+            raise HTTPException(status_code=400, detail="No upload file sent")
+        else:
+            if file.filename == "":
+                raise HTTPException(status_code=400, detail="No selected file")
+            if file:
+                filename = secure_filename(file.filename)
+                folder_path = "SOURCE_DOCUMENTS"
 
-            if not os.path.exists(folder_path):
-                os.makedirs(folder_path)
+                if not os.path.exists(folder_path):
+                    os.makedirs(folder_path)
 
-            file_path = os.path.join(folder_path, filename)
-            file.save(file_path)
+                file_path = os.path.join(folder_path, filename)
+                file.save(file_path)
 
-            return {"response": "File saved successfully"}
+                return {"response": "File saved successfully"}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=e)
+
+@api_app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    while True:
+        data = await websocket.receive_text()
+        await websocket.send_text(f"Message text was: {data}")

requirements.txt

@@ -28,6 +28,7 @@ flask
 requests
 uvicorn
 fastapi
+websockets
 pydantic
 
 # Streamlit related