File size: 6,254 Bytes
25db1d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import os, tempfile
# import pinecone
from pathlib import Path
import traceback
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
from dotenv import load_dotenv
import streamlit as st
load_dotenv()
TMP_DIR = Path(__file__).resolve().parent.joinpath('data', 'tmp')
LOCAL_VECTOR_STORE_DIR = Path(__file__).resolve().parent.joinpath('data', 'vector_store')
# Load environment variables
os.makedirs(TMP_DIR, exist_ok=True)
os.makedirs(LOCAL_VECTOR_STORE_DIR, exist_ok=True)
os.makedirs(TMP_DIR, exist_ok=True)
os.makedirs(LOCAL_VECTOR_STORE_DIR, exist_ok=True)
st.set_page_config(page_title="RAG")
st.title("Retrieval Augmented Generation Engine")
openai_api_key = os.environ.get('OPENAI_API_KEY')
st.session_state.openai_api_key = openai_api_key
def load_documents():
loader = DirectoryLoader(TMP_DIR.as_posix(), glob='**/*.pdf')
documents = loader.load()
return documents
def split_documents(documents):
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
return texts
def embeddings_on_local_vectordb():
# vectordb = Chroma.from_documents(texts, embedding=OpenAIEmbeddings(),
# persist_directory=LOCAL_VECTOR_STORE_DIR.as_posix())
vectordb=Chroma(persist_directory=LOCAL_VECTOR_STORE_DIR.as_posix(), embedding_function=OpenAIEmbeddings())
vectordb.persist()
retriever = vectordb.as_retriever(search_kwargs={'k': 5})
return retriever
# def embeddings_on_pinecone(texts):
# pinecone.init(api_key=st.session_state.pinecone_api_key, environment=st.session_state.pinecone_env)
# embeddings = OpenAIEmbeddings(openai_api_key=st.session_state.openai_api_key)
# vectordb = Pinecone.from_documents(texts, embeddings, index_name=st.session_state.pinecone_index)
# retriever = vectordb.as_retriever()
# return retriever
def query_llm(retriever, query):
try:
qa_chain = ConversationalRetrievalChain.from_llm(
llm=ChatOpenAI(temperature=0, openai_api_key=st.session_state.openai_api_key),
retriever=retriever,
return_source_documents=True,
)
result = qa_chain({'question': query, 'chat_history': st.session_state.messages})
result = result.get('answer')
except Exception as e:
print(f"Exception {e} with traceback : {traceback.format_exc() } occurred for API key: {st.session_state.openai_api_key}")
result = ""
st.session_state.messages.append((query, result))
return result
def input_fields():
#
with st.sidebar:
#
openai_key = st.text_input("OpenAI API key", type="password")
if openai_key != "":
st.session_state.openai_api_key = openai_key
#
# if "pinecone_api_key" in st.secrets:
# st.session_state.pinecone_api_key = st.secrets.pinecone_api_key
# else:
# st.session_state.pinecone_api_key = st.text_input("Pinecone API key", type="password")
#
# if "pinecone_env" in st.secrets:
# st.session_state.pinecone_env = st.secrets.pinecone_env
# else:
# st.session_state.pinecone_env = st.text_input("Pinecone environment")
#
# if "pinecone_index" in st.secrets:
# st.session_state.pinecone_index = st.secrets.pinecone_index
# else:
# st.session_state.pinecone_index = st.text_input("Pinecone index name")
#
# st.session_state.pinecone_db = st.toggle('Use Pinecone Vector DB')
#
st.session_state.source_docs = st.file_uploader(label="Upload Documents", type="pdf", accept_multiple_files=True)
#
retriever = embeddings_on_local_vectordb()
def process_documents():
# if not st.session_state.openai_api_key or not st.session_state.pinecone_api_key or not st.session_state.pinecone_env or not st.session_state.pinecone_index or not st.session_state.source_docs:
if not st.session_state.openai_api_key or not st.session_state.source_docs:
st.warning(f"Please upload the documents and provide the missing fields.")
else:
try:
for source_doc in st.session_state.source_docs:
#
with tempfile.NamedTemporaryFile(delete=False, dir=TMP_DIR.as_posix(), suffix='.pdf') as tmp_file:
tmp_file.write(source_doc.read())
#
documents = load_documents()
#
for _file in TMP_DIR.iterdir():
temp_file = TMP_DIR.joinpath(_file)
temp_file.unlink()
#
texts = split_documents(documents)
print(f"Adding {len(texts)} texts to vector DB")
retriever.add_texts(texts)
retriever.persist()
#
# if not st.session_state.pinecone_db:
# st.session_state.retriever = retriever
# else:
# st.session_state.retriever = embeddings_on_pinecone(texts)
except Exception as e:
st.error(f"An error occurred: {e}")
def boot():
#
input_fields()
#
st.button("Submit Documents", on_click=process_documents)
#
if "messages" not in st.session_state:
st.session_state.messages = []
#
for message in st.session_state.messages:
st.chat_message('human').write(message[0])
st.chat_message('ai').write(message[1])
#
if query := st.chat_input():
st.chat_message("human").write(query)
response = query_llm(retriever, query)
st.chat_message("ai").write(response)
if __name__ == '__main__':
#
boot()
|