Update app.py

app.py

@@ -1,49 +1,22 @@
-import torch
+import logging
+import os
+import shutil
 import subprocess
-import streamlit as st
-from run_localGPT import load_model
-from langchain.vectorstores import Chroma
-from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.chains import RetrievalQA
-from streamlit_extras.add_vertical_space import add_vertical_space
-from langchain.prompts import PromptTemplate
-from langchain.memory import ConversationBufferMemory
-
 
-def model_memory():
-    # Adding history to the model.
-    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.
-
-    {context}
-
-    {history}
-    Question: {question}
-    Helpful Answer:"""
-
-    prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)
-    memory = ConversationBufferMemory(input_key="question", memory_key="history")
-
-    return prompt, memory
+import torch
+from flask import Flask, jsonify, request
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceInstructEmbeddings
 
+# from langchain.embeddings import HuggingFaceEmbeddings
+from run_localGPT import load_model
+from prompt_template_utils import get_prompt_template
 
-# Sidebar contents
-with st.sidebar:
-    st.title("🤗💬 Converse with your Data")
-    st.markdown(
-        """
-    ## About
-    This app is an LLM-powered chatbot built using:
-    - [Streamlit](https://streamlit.io/)
-    - [LangChain](https://python.langchain.com/)
-    - [LocalGPT](https://github.com/PromtEngineer/localGPT) 
- 
-    """
-    )
-    add_vertical_space(5)
-    st.write("Made with ❤️ by [Prompt Engineer](https://youtube.com/@engineerprompt)")
+# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.vectorstores import Chroma
+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"
@@ -52,71 +25,160 @@ elif torch.cuda.is_available():
 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})
+
+# uncomment the following line if you used HuggingFaceEmbeddings in the ingest.py
+# EMBEDDINGS = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+# if os.path.exists(PERSIST_DIRECTORY):
+#     try:
+#         shutil.rmtree(PERSIST_DIRECTORY)
+#     except OSError as e:
+#         print(f"Error: {e.filename} - {e.strerror}.")
+# else:
+#     print("The directory does not exist")
 
-# if "result" not in st.session_state:
-#     # Run the document ingestion process.
-#     run_langest_commands = ["python", "ingest.py"]
+# run_langest_commands = ["python", "ingest.py"]
+# if DEVICE_TYPE == "cpu":
 #     run_langest_commands.append("--device_type")
 #     run_langest_commands.append(DEVICE_TYPE)
 
-#     result = subprocess.run(run_langest_commands, capture_output=True)
-#     st.session_state.result = result
+# result = subprocess.run(run_langest_commands, capture_output=True)
+# if result.returncode != 0:
+#     raise FileNotFoundError(
+#         "No files were found inside SOURCE_DOCUMENTS, please put a starter file inside before starting the API!"
+#     )
 
-# Define the retreiver
 # load the vectorstore
-if "EMBEDDINGS" not in st.session_state:
-    EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": DEVICE_TYPE})
-    st.session_state.EMBEDDINGS = EMBEDDINGS
-
-if "DB" not in st.session_state:
-    DB = Chroma(
-        persist_directory=PERSIST_DIRECTORY,
-        embedding_function=st.session_state.EMBEDDINGS,
-        client_settings=CHROMA_SETTINGS,
-    )
-    st.session_state.DB = DB
-
-if "RETRIEVER" not in st.session_state:
-    RETRIEVER = DB.as_retriever()
-    st.session_state.RETRIEVER = RETRIEVER
-
-if "LLM" not in st.session_state:
-    LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
-    st.session_state["LLM"] = LLM
-
-
-if "QA" not in st.session_state:
-    prompt, memory = model_memory()
-
-    QA = RetrievalQA.from_chain_type(
-        llm=LLM,
-        chain_type="stuff",
-        retriever=RETRIEVER,
-        return_source_documents=True,
-        chain_type_kwargs={"prompt": prompt, "memory": memory},
+DB = Chroma(
+    persist_directory=PERSIST_DIRECTORY,
+    embedding_function=EMBEDDINGS,
+    client_settings=CHROMA_SETTINGS,
+)
+
+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)
+
+QA = RetrievalQA.from_chain_type(
+    llm=LLM,
+    chain_type="stuff",
+    retriever=RETRIEVER,
+    return_source_documents=SHOW_SOURCES,
+    chain_type_kwargs={
+        "prompt": prompt,
+    },
+)
+
+app = Flask(__name__)
+
+
+@app.route("/api/delete_source", methods=["GET"])
+def delete_source_route():
+    folder_name = "SOURCE_DOCUMENTS"
+
+    if os.path.exists(folder_name):
+        shutil.rmtree(folder_name)
+
+    os.makedirs(folder_name)
+
+    return jsonify({"message": f"Folder '{folder_name}' successfully deleted and recreated."})
+
+
+@app.route("/api/save_document", methods=["GET", "POST"])
+def save_document_route():
+    if "document" not in request.files:
+        return "No document part", 400
+    file = request.files["document"]
+    if file.filename == "":
+        return "No selected file", 400
+    if file:
+        filename = secure_filename(file.filename)
+        folder_path = "SOURCE_DOCUMENTS"
+        if not os.path.exists(folder_path):
+            os.makedirs(folder_path)
+        file_path = os.path.join(folder_path, filename)
+        file.save(file_path)
+        return "File saved successfully", 200
+
+
+@app.route("/api/run_ingest", methods=["GET"])
+def run_ingest_route():
+    global DB
+    global RETRIEVER
+    global QA
+    try:
+        if os.path.exists(PERSIST_DIRECTORY):
+            try:
+                shutil.rmtree(PERSIST_DIRECTORY)
+            except OSError as e:
+                print(f"Error: {e.filename} - {e.strerror}.")
+        else:
+            print("The directory does not exist")
+
+        run_langest_commands = ["python", "ingest.py"]
+        if DEVICE_TYPE == "cpu":
+            run_langest_commands.append("--device_type")
+            run_langest_commands.append(DEVICE_TYPE)
+
+        result = subprocess.run(run_langest_commands, capture_output=True)
+        if result.returncode != 0:
+            return "Script execution failed: {}".format(result.stderr.decode("utf-8")), 500
+        # load the vectorstore
+        DB = Chroma(
+            persist_directory=PERSIST_DIRECTORY,
+            embedding_function=EMBEDDINGS,
+            client_settings=CHROMA_SETTINGS,
+        )
+        RETRIEVER = DB.as_retriever()
+        prompt, memory = get_prompt_template(promptTemplate_type="llama", history=False)
+
+        QA = RetrievalQA.from_chain_type(
+            llm=LLM,
+            chain_type="stuff",
+            retriever=RETRIEVER,
+            return_source_documents=SHOW_SOURCES,
+            chain_type_kwargs={
+                "prompt": prompt,
+            },
+        )
+        return "Script executed successfully: {}".format(result.stdout.decode("utf-8")), 200
+    except Exception as e:
+        return f"Error occurred: {str(e)}", 500
+
+
+@app.route("/api/prompt_route", methods=["GET", "POST"])
+def prompt_route():
+    global QA
+    user_prompt = request.form.get("user_prompt")
+    if user_prompt:
+        # print(f'User Prompt: {user_prompt}')
+        # Get the answer from the chain
+        res = QA(user_prompt)
+        answer, docs = res["result"], res["source_documents"]
+
+        prompt_response_dict = {
+            "Prompt": user_prompt,
+            "Answer": answer,
+        }
+
+        prompt_response_dict["Sources"] = []
+        for document in docs:
+            prompt_response_dict["Sources"].append(
+                (os.path.basename(str(document.metadata["source"])), str(document.page_content))
+            )
+
+        return jsonify(prompt_response_dict), 200
+    else:
+        return "No user prompt received", 400
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO
     )
-    st.session_state["QA"] = QA
-
-st.title("LocalGPT App 💬")
-# Create a text input box for the user
-prompt = st.text_input("Input your prompt here")
-# while True:
-
-# If the user hits enter
-if prompt:
-    # Then pass the prompt to the LLM
-    response = st.session_state["QA"](prompt)
-    answer, docs = response["result"], response["source_documents"]
-    # ...and write it out to the screen
-    st.write(answer)
-
-    # With a streamlit expander
-    with st.expander("Document Similarity Search"):
-        # Find the relevant pages
-        search = st.session_state.DB.similarity_search_with_score(prompt)
-        # Write out the first
-        for i, doc in enumerate(search):
-            # print(doc)
-            st.write(f"Source Document # {i+1} : {doc[0].metadata['source'].split('/')[-1]}")
-            st.write(doc[0].page_content)
-            st.write("--------------------------------")
+    app.run(debug=False, port=5110)