Create app.py

app.py

@@ -0,0 +1,157 @@
+from flask import Flask, render_template, jsonify, request, redirect, url_for
+from flask_wtf.csrf import CSRFProtect
+
+# from tavily import TavilyClient
+
+from dotenv import load_dotenv
+import os
+
+from langchain_community.document_loaders import TextLoader
+from langchain_community.vectorstores import Chroma
+from langchain_text_splitters import CharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_openai import OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains import create_retrieval_chain
+from langchain import hub
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.chains.question_answering import load_qa_chain
+from langchain.prompts import PromptTemplate
+
+import time
+load_dotenv()
+
+# TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+# tavily = TavilyClient(api_key=TAVILY_API_KEY)
+
+app = Flask(__name__, static_folder='static')
+app.config['SECRET_KEY'] = 'secret'
+csrf = CSRFProtect(app)
+
+text_splitter = CharacterTextSplitter(separator = "\n", chunk_size=1000, chunk_overlap=200, length_function = len)
+embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
+retrieval_qa_chat_prompt = hub.pull("langchain-ai/retrieval-qa-chat")
+llm = ChatOpenAI(api_key=OPENAI_API_KEY)
+
+vectordb_path = "./vector_db"
+
+@app.route('/')
+def home():
+    return redirect(url_for('search_view'))
+
+@app.route('/search_view')
+def search_view():
+    return render_template('search.html')
+
+@app.route('/rag_view')
+def rag_view():
+    dbs = [f.name for f in os.scandir(vectordb_path) if f.is_dir()]
+    return render_template('rag.html', dbs = dbs)
+
+@app.route('/query', methods=['POST'])
+def query():
+    if request.method == "POST":
+        prompt = request.get_json().get("prompt")
+        title = request.get_json().get("title")
+        db = request.get_json().get("db")
+        
+        # if title == "search":
+        #     response = tavily.search(query=prompt, include_images=True, include_answer=True, max_results=5)
+            
+        #     output = response['answer'] + "\n"
+        #     for res in response['results']:
+        #         output += f"\nTitle: {res['title']}\nURL: {res['url']}\nContent: {res['content']}\n"
+        
+        #     data = {"success": "ok", "response": output, "images": response['images']}
+
+        #     return jsonify(data)    
+
+        if title == "rag":
+            if db != "":
+                template = """Please answer to human's input based on context. If the input is not mentioned in context, output something like 'I don't know'.
+                Context: {context}
+                Human: {human_input}
+                Your Response as Chatbot:"""
+                
+                prompt_s = PromptTemplate(
+                    input_variables=["human_input", "context"],
+                    template=template
+                )
+
+                db = Chroma(persist_directory=os.path.join(vectordb_path, db), embedding_function=embeddings)
+                
+                docs = db.similarity_search(prompt)
+                
+                llm = ChatOpenAI(model="gpt-4-1106-preview", api_key=OPENAI_API_KEY)
+
+                stuff_chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_s)
+                output = stuff_chain({"input_documents": docs, "human_input": prompt}, return_only_outputs=False)
+
+                final_answer = output["output_text"]
+                # prompt = ChatPromptTemplate.from_messages(
+                #     [("system", "Please answer to user's query based on following context.\n\nContext: {context}")]
+                # )
+                
+                
+                # chain = create_stuff_documents_chain(llm, prompt)
+
+                # answer = chain.invoke({"context": docs, "prompt": prompt})
+
+                data = {"success": "ok", "response": final_answer}
+                
+                return jsonify(data)
+            else:
+                data = {"success": "ok", "response": "Please select database."}
+
+                return jsonify(data)
+        
+@app.route('/uploadDocuments', methods=['POST'])
+@csrf.exempt
+def uploadDocuments():
+    # uploaded_files = request.files.getlist('files[]')
+    dbname = request.form.get('dbname')
+    uploaded_files = ['https://www.airbus.com/sites/g/files/jlcbta136/files/2024-03/Airbus-Annual-Report-2023.pdf', 'https://www.singaporeair.com/saar5/pdf/Investor-Relations/Annual-Report/annualreport2223.pdf']
+    if dbname == "":
+        return {"success": "db"}
+    
+    if len(uploaded_files) > 0:
+        for file in uploaded_files:
+            file.save(f"uploads/{file.filename}")
+        
+            if file.filename.endswith(".txt"):
+                loader = TextLoader(f"uploads/{file.filename}", encoding='utf-8')
+            else:
+                loader = PyPDFLoader(f"uploads/{file.filename}")
+                
+            data = loader.load()
+            texts = text_splitter.split_documents(data)
+            
+            Chroma.from_documents(texts, embeddings, persist_directory=os.path.join(vectordb_path, dbname))
+            
+        return {'success': "ok"}
+
+    else:
+        return {"success": "bad"}
+
+@app.route('/dbcreate', methods=['POST'])
+@csrf.exempt
+def dbcreate():
+    dbname = request.get_json().get("dbname")
+    
+    if not os.path.exists(os.path.join(vectordb_path, dbname)):
+        os.makedirs(os.path.join(vectordb_path, dbname))
+        return {'success': "ok"}
+    else:
+        return {'success': 'bad'}
+    
+# import gradio as gr
+# chatbot = gr.Chatbot(avatar_images=["user.png", "bot.jpg"], height=600)
+# clear_but = gr.Button(value="Clear Chat")
+# demo = gr.ChatInterface(fn=search, title="Mediate.com Chatbot Prototype", multimodal=False, retry_btn=None, undo_btn=None, clear_btn=clear_but, chatbot=chatbot)
+
+
+if __name__ == '__main__':
+    app.run(debug=True)