import os
import tempfile
import json
import streamlit as st
from streamlit_chat import message
from preprocessing import Model
from io import BytesIO
import pickle
# Home Page Setup
st.set_page_config(
page_title="PDF Insight Pro",
page_icon="📄",
layout="centered",
)
# Custom CSS for a more polished look
st.markdown("""
""", unsafe_allow_html=True)
# Custom title and header
st.title("📄 PDF Insight Pro")
st.subheader("Empower Your Documents with AI-Driven Insights")
def display_messages():
"""
Displays the chat messages in the Streamlit UI.
"""
st.subheader("🗨️ Conversation")
st.markdown("---")
for i, (msg, is_user) in enumerate(st.session_state["messages"]):
message(msg, is_user=is_user, key=str(i))
st.session_state["process_input_spinner"] = st.empty()
def process_user_input():
"""
Processes the user input by generating a response from the assistant.
"""
if st.session_state["user_input"] and len(st.session_state["user_input"].strip()) > 0:
user_input = st.session_state["user_input"].strip()
with st.session_state["process_input_spinner"], st.spinner("Analyzing..."):
agent_response = st.session_state["assistant"].get_response(
user_input,
st.session_state["temperature"],
st.session_state["max_tokens"],
st.session_state["model"]
)
st.session_state["messages"].append((user_input, True))
st.session_state["messages"].append((agent_response, False))
st.session_state["user_input"] = ""
# Save chat history temporarily on local storage
with open("chat_history.pkl", "wb") as f:
pickle.dump(st.session_state["messages"], f)
def process_file():
"""
Processes the uploaded PDF file and appends its content to the context.
"""
for file in st.session_state["file_uploader"]:
with tempfile.NamedTemporaryFile(delete=False) as tf:
tf.write(file.getbuffer())
file_path = tf.name
with st.session_state["process_file_spinner"], st.spinner(f"Processing {file.name}..."):
try:
st.session_state["assistant"].add_to_context(file_path)
except Exception as e:
st.error(f"Failed to process file {file.name}: {str(e)}")
os.remove(file_path)
def download_chat_history():
"""
Allows users to download chat history in HTML or JSON format.
"""
# Convert messages to JSON format
chat_data = [{"role": "user" if is_user else "assistant", "content": msg} for msg, is_user in st.session_state["messages"]]
# Download as JSON
json_data = json.dumps(chat_data, indent=4)
st.download_button(
label="💾 Download Chat History as JSON",
data=json_data,
file_name="chat_history.json",
mime="application/json"
)
# Download as HTML
html_data = "Chat History
"
for entry in chat_data:
role = "User" if entry["role"] == "user" else "Assistant"
html_data += f"- {role}: {entry['content']}
"
html_data += "
"
st.download_button(
label="💾 Download Chat History as HTML",
data=html_data,
file_name="chat_history.html",
mime="text/html"
)
def main_page():
"""
Main function to set up the Streamlit UI and handle user interactions.
"""
# Initialize session state variables
if "messages" not in st.session_state:
st.session_state["messages"] = []
if "assistant" not in st.session_state:
st.session_state["assistant"] = Model()
if "user_input" not in st.session_state:
st.session_state["user_input"] = ""
if "temperature" not in st.session_state:
st.session_state["temperature"] = 0.5
if "max_tokens" not in st.session_state:
st.session_state["max_tokens"] = 550
if "model" not in st.session_state:
st.session_state["model"] = "llama-3.1-8b-instant"
# File uploader
st.subheader("📤 Upload Your PDF Documents")
st.file_uploader(
"Choose PDF files to analyze",
type=["pdf"],
key="file_uploader",
on_change=process_file,
accept_multiple_files=True,
)
st.session_state["process_file_spinner"] = st.empty()
# Document management section
if st.session_state["assistant"].contexts:
st.subheader("🗂️ Manage Uploaded Documents")
for i, context in enumerate(st.session_state["assistant"].contexts):
st.text_area(f"Document {i+1} Context", context[:500] + "..." if len(context) > 500 else context, height=100)
if st.button(f"Remove Document {i+1}"):
st.session_state["assistant"].remove_from_context(i)
# Model settings
with st.expander("⚙️ Customize AI Settings", expanded=True):
st.slider("Sampling Temperature", min_value=0.0, max_value=1.0, step=0.1, key="temperature", help="Higher values make output more random.")
st.slider("Max Tokens", min_value=750, max_value=5000, step=50, key="max_tokens", help="Limits the length of the response.")
st.selectbox("Choose AI Model", ["llama-3.1-8b-instant", "llama3-70b-8192", "gemma-7b-it"], key="model")
# Display messages and input box
display_messages()
st.text_input("Type your query and hit Enter", key="user_input", on_change=process_user_input, placeholder="Ask something about your documents...")
# Download chat history section
st.subheader("💾 Download Chat History")
download_chat_history()
# Developer info and bug report
st.subheader("🐞 Bug Report")
st.markdown("""
If you encounter any bugs or issues while using the app, please send a bug report to the developer. You can include a screenshot (optional) to help identify the problem.\n
""")
st.subheader("💡 Suggestions")
st.markdown("""
Suggestions to improve the app's UI and user interface are also welcome. Feel free to reach out to the developer with your suggestions.\n
""")
st.subheader("👨💻 Developer Info")
st.markdown("""
**Developer**: Jatin Mehra\n
**Email**: jatinmehra119@gmail.com\n
**Mobile**: 9910364780\n
""")
if __name__ == "__main__":
main_page()