Triayushi/UI-RAG-pdf-retriever-streamlit
0
1# import altair as alt2# import numpy as np3# import pandas as pd4# import streamlit as st5 6# """7# # Welcome to Streamlit!8 9# Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.10# If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community11# forums](https://discuss.streamlit.io).12 13# In the meantime, below is an example of what you can do with just a few lines of code:14# """15 16# num_points = st.slider("Number of points in spiral", 1, 10000, 1100)17# num_turns = st.slider("Number of turns in spiral", 1, 300, 31)18 19# indices = np.linspace(0, 1, num_points)20# theta = 2 * np.pi * num_turns * indices21# radius = indices22 23# x = radius * np.cos(theta)24# y = radius * np.sin(theta)25 26# df = pd.DataFrame({27# "x": x,28# "y": y,29# "idx": indices,30# "rand": np.random.randn(num_points),31# })32 33# st.altair_chart(alt.Chart(df, height=700, width=700)34# .mark_point(filled=True)35# .encode(36# x=alt.X("x", axis=None),37# y=alt.Y("y", axis=None),38# color=alt.Color("idx", legend=None, scale=alt.Scale()),39# size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),40# ))41 42 43import os44import boto345import tempfile46from langchain_community.document_loaders import PyPDFLoader47from langchain.text_splitter import RecursiveCharacterTextSplitter48from langchain_community.vectorstores import Chroma49from langchain_aws import BedrockLLM, BedrockEmbeddings50from langchain.chains import RetrievalQA51import streamlit as st52 53# AWS Configuration54AWS_REGION = os.getenv("AWS_REGION")55AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")56AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")57 58# Initialize Bedrock client59bedrock_client = boto3.client(60 service_name="bedrock-runtime",61 region_name=AWS_REGION,62 aws_access_key_id=AWS_ACCESS_KEY_ID,63 aws_secret_access_key=AWS_SECRET_ACCESS_KEY64)65 66# Streamlit page config67st.set_page_config(page_title="PDF Chatbot", layout="wide")68st.title("PDF Q&A Chatbot")69 70# ChromaDB path71CHROMA_PATH = os.path.join(os.getcwd(), "chroma_db")72 73# Initialize session state74if "chat_history" not in st.session_state:75 st.session_state.chat_history = []76if "vectorstore" not in st.session_state:77 st.session_state.vectorstore = None78if "pdf_processed" not in st.session_state:79 st.session_state.pdf_processed = False80 81def clear_vectorstore():82 """Clear existing ChromaDB data."""83 if os.path.exists(CHROMA_PATH):84 try:85 Chroma(persist_directory=CHROMA_PATH, embedding_function=None).delete_collection()86 st.info("Cleared previous vector store.")87 except Exception as e:88 st.warning(f"Failed to clear previous vector store: {str(e)}")89 st.session_state.vectorstore = None90 st.session_state.pdf_processed = False91 92def check_directory_permissions(path):93 """Check if directory is writable."""94 try:95 test_file = os.path.join(path, "test.txt")96 with open(test_file, "w") as f:97 f.write("test")98 os.remove(test_file)99 return True100 except Exception as e:101 st.error(f"Directory {path} is not writable: {str(e)}")102 return False103 104def main():105 # Check ChromaDB directory permissions106 if not check_directory_permissions(CHROMA_PATH):107 st.error("Cannot proceed due to ChromaDB directory permission issues.")108 return109 110 # Initialize embeddings111 try:112 embeddings = BedrockEmbeddings(113 client=bedrock_client,114 model_id="amazon.titan-embed-text-v1"115 )116 except Exception as e:117 st.error(f"Failed to initialize Bedrock embeddings: {str(e)}")118 return119 120 # Sidebar for PDF upload and processing121 with st.sidebar:122 st.title("Upload PDF")123 with st.form(key="pdf_upload_form"):124 uploaded_file = st.file_uploader("Choose a PDF file", type="pdf", key="pdf_uploader")125 submit_button = st.form_submit_button("Process PDF")126 if submit_button and uploaded_file:127 with st.spinner("Processing..."):128 try:129 st.write(f"Uploading file: {uploaded_file.name}, size: {uploaded_file.size} bytes")130 # Clear previous vector store131 clear_vectorstore()132 133 # Save uploaded PDF to temporary file134 with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:135 tmp_file.write(uploaded_file.getvalue())136 tmp_file_path = tmp_file.name137 138 # Load PDF139 loader = PyPDFLoader(tmp_file_path)140 pages = loader.load()141 st.write(f"Loaded {len(pages)} pages")142 143 # Split documents144 text_splitter = RecursiveCharacterTextSplitter(145 chunk_size=1000,146 chunk_overlap=200147 )148 chunks = text_splitter.split_documents(pages)149 st.write(f"Split into {len(chunks)} chunks")150 151 # Initialize ChromaDB152 st.session_state.vectorstore = Chroma(153 persist_directory=CHROMA_PATH,154 embedding_function=embeddings155 )156 st.session_state.vectorstore.add_documents(chunks)157 st.session_state.vectorstore.persist()158 st.session_state.pdf_processed = True159 160 # Clean up temporary file161 os.unlink(tmp_file_path)162 st.success(f"Processed {uploaded_file.name} successfully!")163 except Exception as e:164 st.error(f"Error processing PDF: {str(e)}")165 st.session_state.pdf_processed = False166 167 # Display chat history168 st.header("Chat with PDF")169 for message in st.session_state.chat_history:170 with st.chat_message(message["role"]):171 st.write(message["content"])172 173 # User input174 user_question = st.chat_input("Ask about the uploaded PDF:")175 176 if user_question:177 if not st.session_state.pdf_processed or not st.session_state.vectorstore:178 st.warning("Please upload and process a PDF first.")179 return180 181 try:182 # Add user question to chat history183 st.session_state.chat_history.append({"role": "user", "content": user_question})184 with st.chat_message("user"):185 st.write(user_question)186 187 # Initialize retriever188 retriever = st.session_state.vectorstore.as_retriever(search_kwargs={'k': 5})189 190 # Check if retriever has documents191 if not retriever.invoke(user_question):192 st.warning("No relevant content found in the PDF for this question.")193 return194 195 # Prepare conversation history for Claude196 messages = [197 {"role": "user" if msg["role"] == "user" else "assistant", "content": msg["content"]}198 for msg in st.session_state.chat_history199 ]200 201 # Combine user question with context instruction202 contextual_query = f"Answer the following question based on the content of the uploaded PDF: {user_question}"203 204 # Initialize Bedrock LLM205 llm = BedrockLLM(206 client=bedrock_client,207 model_id="amazon.titan-embed-text-v1",208 model_kwargs={209 "max_tokens": 1000,210 "anthropic_version": "bedrock-2023-05-31",211 "temperature": 0212 }213 )214 215 # Create RAG chain216 chain = RetrievalQA.from_chain_type(217 llm=llm,218 chain_type="stuff",219 retriever=retriever220 )221 222 # Get answer223 answer = chain.invoke({"query": contextual_query, "messages": messages})224 response = answer['result']225 226 # Add assistant response to chat history227 st.session_state.chat_history.append({"role": "assistant", "content": response})228 with st.chat_message("assistant"):229 st.write(response)230 231 except Exception as e:232 st.error(f"Error generating answer: {str(e)}")233 234if __name__ == "__main__":235 main()