mirror of
https://github.com/arc53/DocsGPT
synced 2024-11-02 03:40:17 +00:00
142 lines
4.2 KiB
Python
142 lines
4.2 KiB
Python
import os
|
|
import pickle
|
|
import dotenv
|
|
import datetime
|
|
from flask import Flask, request, render_template
|
|
# os.environ["LANGCHAIN_HANDLER"] = "langchain"
|
|
import faiss
|
|
from langchain import FAISS
|
|
from langchain import OpenAI, VectorDBQA, HuggingFaceHub, Cohere
|
|
from langchain.chains.question_answering import load_qa_chain
|
|
from langchain.prompts import PromptTemplate
|
|
import requests
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
|
|
# from manifest import Manifest
|
|
# from langchain.llms.manifest import ManifestWrapper
|
|
|
|
# manifest = Manifest(
|
|
# client_name = "huggingface",
|
|
# client_connection = "http://127.0.0.1:5000"
|
|
# )
|
|
|
|
|
|
# Redirect PosixPath to WindowsPath on Windows
|
|
import platform
|
|
if platform.system() == "Windows":
|
|
import pathlib
|
|
temp = pathlib.PosixPath
|
|
pathlib.PosixPath = pathlib.WindowsPath
|
|
|
|
# loading the .env file
|
|
dotenv.load_dotenv()
|
|
|
|
|
|
with open("combine_prompt.txt", "r") as f:
|
|
template = f.read()
|
|
|
|
# check if OPENAI_API_KEY is set
|
|
if os.getenv("OPENAI_API_KEY") is not None:
|
|
api_key_set = True
|
|
|
|
else:
|
|
api_key_set = False
|
|
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
@app.route("/")
|
|
def home():
|
|
return render_template("index.html", api_key_set=api_key_set)
|
|
|
|
|
|
@app.route("/api/answer", methods=["POST"])
|
|
def api_answer():
|
|
data = request.get_json()
|
|
question = data["question"]
|
|
if not api_key_set:
|
|
api_key = data["api_key"]
|
|
else:
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
|
|
# check if the vectorstore is set
|
|
if "active_docs" in data:
|
|
vectorstore = "vectors/" + data["active_docs"]
|
|
if data['active_docs'] == "default":
|
|
vectorstore = ""
|
|
else:
|
|
vectorstore = ""
|
|
|
|
|
|
# loading the index and the store and the prompt template
|
|
docsearch = FAISS.load_local(vectorstore, OpenAIEmbeddings(openai_api_key=api_key))
|
|
|
|
# create a prompt template
|
|
c_prompt = PromptTemplate(input_variables=["summaries", "question"], template=template)
|
|
# create a chain with the prompt template and the store
|
|
|
|
#llm = ManifestWrapper(client=manifest, llm_kwargs={"temperature": 0.001, "max_tokens": 2048})
|
|
llm = OpenAI(openai_api_key=api_key, temperature=0)
|
|
#llm = HuggingFaceHub(repo_id="bigscience/bloom", huggingfacehub_api_token=api_key)
|
|
# llm = Cohere(model="command-xlarge-nightly", cohere_api_key=api_key)
|
|
|
|
qa_chain = load_qa_chain(llm = llm, chain_type="map_reduce",
|
|
combine_prompt=c_prompt)
|
|
|
|
|
|
chain = VectorDBQA(combine_documents_chain=qa_chain, vectorstore=docsearch, k=2)
|
|
|
|
# fetch the answer
|
|
result = chain({"query": question})
|
|
print(result)
|
|
|
|
# some formatting for the frontend
|
|
result['answer'] = result['result']
|
|
result['answer'] = result['answer'].replace("\\n", "<br>")
|
|
result['answer'] = result['answer'].replace("SOURCES:", "")
|
|
# mock result
|
|
# result = {
|
|
# "answer": "The answer is 42",
|
|
# "sources": ["https://en.wikipedia.org/wiki/42_(number)", "https://en.wikipedia.org/wiki/42_(number)"]
|
|
# }
|
|
return result
|
|
|
|
@app.route("/api/docs_check", methods=["POST"])
|
|
def check_docs():
|
|
# check if docs exist in a vectorstore folder
|
|
data = request.get_json()
|
|
vectorstore = "vectors/" + data["docs"]
|
|
base_path = 'https://raw.githubusercontent.com/arc53/DocsHUB/main/'
|
|
#
|
|
if os.path.exists(vectorstore):
|
|
return {"status": 'exists'}
|
|
else:
|
|
r = requests.get(base_path + vectorstore + "docs.index")
|
|
# save to vectors directory
|
|
# check if the directory exists
|
|
if not os.path.exists(vectorstore):
|
|
os.makedirs(vectorstore)
|
|
|
|
with open(vectorstore + "docs.index", "wb") as f:
|
|
f.write(r.content)
|
|
# download the store
|
|
r = requests.get(base_path + vectorstore + "faiss_store.pkl")
|
|
with open(vectorstore + "faiss_store.pkl", "wb") as f:
|
|
f.write(r.content)
|
|
|
|
return {"status": 'loaded'}
|
|
|
|
# handling CORS
|
|
@app.after_request
|
|
def after_request(response):
|
|
response.headers.add('Access-Control-Allow-Origin', '*')
|
|
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
|
response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
|
|
return response
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|