File size: 5,377 Bytes
2639832 |
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 |
import json
import os
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.vectorstores import FAISS
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
import streamlit as st
from dotenv import load_dotenv
from langchain_core.runnables.history import RunnableWithMessageHistory
load_dotenv()
# Langsmith Tracking
os.environ['LANGCHAIN_API_KEY'] = os.getenv('LANGCHAIN_API_KEY')
os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_PROJECT'] = "Rag with chat history"
os.environ['GROQ_API_KEY'] = os.getenv('GROQ_API_KEY')
os.environ["HF_TOKEN"] = os.getenv('HF_TOKEN')
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Initialize embeddings
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
# Load or initialize sessions.json
session_file = 'sessions.json'
if not os.path.exists(session_file):
with open(session_file, 'w') as f:
json.dump({"current_session_id": 1}, f)
def get_new_session_id():
with open(session_file, 'r+') as f:
data = json.load(f)
session_id = data['current_session_id']
data['current_session_id'] += 1
f.seek(0)
json.dump(data, f)
f.truncate()
return session_id
# Set up Streamlit App
st.title("Rag with chat history")
llm = ChatGroq(model="llama-3.1-70b-Versatile")
# Get or create session ID
if 'session_id' not in st.session_state:
st.session_state.session_id = get_new_session_id()
session_id = st.session_state.session_id
st.write(f"Session ID: {session_id}")
# Statefully manage chat history
if 'store' not in st.session_state:
st.session_state.store = {}
uploaded_files = st.file_uploader("Choose a PDF file", type="pdf", accept_multiple_files=True)
# Process Uploaded Files:
if uploaded_files:
documents = []
for uploaded_file in uploaded_files:
temppdf = f"./temp.pdf"
with open(temppdf, "wb") as file:
file.write(uploaded_file.getvalue())
loader = PyPDFLoader(temppdf)
docs = loader.load()
documents.extend(docs)
# Split and create embedding documents
text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500)
splits = text_splitter.split_documents(documents)
vector_store = FAISS.from_documents(documents=splits, embedding=embeddings)
retriever = vector_store.as_retriever()
contextualize_q_systemprompt = (
"Given a chat history and the latest user question "
"which might reference context in the chat history, "
"formulate a standalone question which can be understood "
"without the chat history. Do not answer the question, "
"just reformulate it if needed and otherwise return it as it is."
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
[
("system", contextualize_q_systemprompt),
MessagesPlaceholder("chat_history"),
("human", "{input}")
]
)
history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
# Answer question prompt
system_prompt = (
"You are an assistant for question-answering tasks. "
"Use the following pieces of retrieved context to answer the question. "
"If you don't have enough context, you can say that you "
"don't know. Use three sentences maximum and keep the "
"answer concise."
"\n\n"
"{context}"
)
qa_prompt = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
MessagesPlaceholder("chat_history"),
("human", "{input}")
]
)
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
def get_session_history(session: str) -> BaseChatMessageHistory:
if session_id not in st.session_state.store:
st.session_state.store[session_id] = ChatMessageHistory()
return st.session_state.store[session_id]
conversational_rag_chain = RunnableWithMessageHistory(
rag_chain, get_session_history,
input_messages_key="input",
history_messages_key="chat_history",
output_messages_key="answer"
)
user_input = st.text_input("Ask a question")
if user_input:
session_history = get_session_history(session_id)
response = conversational_rag_chain.invoke(
{"input": user_input},
config={"configurable": {"session_id": session_id}},
)
st.session_state.store[session_id] = session_history
st.write(st.session_state.store)
st.write("Assistant:", response["answer"])
st.write("Chat History:", session_history.messages)
else:
st.write("Please upload a file") |