Spaces:
Sleeping
Sleeping
| # ---- Standard libs ---- | |
| import os, tempfile, shutil, uuid | |
| from pathlib import Path | |
| from datetime import datetime | |
| import streamlit as st | |
| from dotenv import load_dotenv | |
| # ---- LangChain stack ---- | |
| from langchain_groq import ChatGroq | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| from langchain_community.vectorstores import Chroma | |
| from langchain.chains import RetrievalQA | |
| from langchain_community.document_loaders import UnstructuredFileLoader | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| # ---- Load .env ---- | |
| load_dotenv() | |
| # ---- Session Management ---- | |
| if "session_id" not in st.session_state: | |
| st.session_state.session_id = str(uuid.uuid4()) | |
| def cleanup_session_data(): | |
| session_path = Path(f"./rag_db/{st.session_state.session_id}") | |
| if session_path.exists(): | |
| try: | |
| shutil.rmtree(session_path) | |
| except Exception as e: | |
| print(f"⚠️ Could not delete {session_path}: {e}") | |
| if "cleaned_up" not in st.session_state: | |
| cleanup_session_data() | |
| st.session_state["cleaned_up"] = True | |
| # ---- Initialize RAG Chain ---- | |
| def init_rag_chain(): | |
| embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") | |
| vectordb = Chroma( | |
| persist_directory=f"./rag_db/{st.session_state.session_id}", | |
| embedding_function=embedding | |
| ) | |
| llm = ChatGroq( | |
| groq_api_key=os.getenv("GROQ_API_KEY"), | |
| model="llama3-70b-8192", | |
| temperature=0.2, | |
| ) | |
| return RetrievalQA.from_chain_type( | |
| llm=llm, | |
| retriever=vectordb.as_retriever(search_kwargs={"k": 4}), | |
| return_source_documents=False | |
| ) | |
| qa_chain = init_rag_chain() | |
| def ingest_pdf(file_bytes: bytes, file_name: str): | |
| embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") | |
| vectordb = Chroma( | |
| persist_directory=f"./rag_db/{st.session_state.session_id}", | |
| embedding_function=embedding | |
| ) | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: | |
| tmp.write(file_bytes) | |
| tmp_path = tmp.name | |
| docs = UnstructuredFileLoader(tmp_path).load() | |
| splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) | |
| chunks = splitter.split_documents(docs) | |
| if chunks: | |
| vectordb.add_documents(chunks) | |
| vectordb.persist() | |
| # ---- Streamlit Config ---- | |
| st.set_page_config( | |
| page_title="📄 Multimodal Chat Assistant", | |
| page_icon="🤖", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # ---- Custom CSS ---- | |
| st.markdown(""" | |
| <style> | |
| .block-container { | |
| padding-top: 2rem; | |
| padding-bottom: 2rem; | |
| max-width: 800px; | |
| margin: auto; | |
| } | |
| .title { | |
| font-size: 3rem; | |
| font-weight: bold; | |
| color: #4A90E2; | |
| text-align: center; | |
| margin-bottom: 1.5rem; | |
| } | |
| .main-container { | |
| max-width: 800px; | |
| margin: 0 auto; | |
| padding: 1rem; | |
| } | |
| .chat-history { | |
| max-height: 450px; | |
| overflow-y: auto; | |
| background-color: #F9F9F9; | |
| border-radius: 12px; | |
| padding: 1rem; | |
| border: 1px solid #DDD; | |
| margin-bottom: 1rem; | |
| min-height: 100px; | |
| } | |
| .message-user { | |
| background-color: #DDF3FF; | |
| color: #000; | |
| padding: 0.8rem; | |
| border-radius: 12px; | |
| margin-bottom: 0.5rem; | |
| text-align: left; | |
| word-wrap: break-word; | |
| } | |
| .message-bot { | |
| background-color: #E3FFE1; | |
| color: #000; | |
| padding: 0.8rem; | |
| border-radius: 12px; | |
| margin-bottom: 0.5rem; | |
| text-align: left; | |
| word-wrap: break-word; | |
| } | |
| .stTextInput > div > div > input { | |
| font-size: 1rem; | |
| padding: 0.75rem; | |
| border-radius: 8px; | |
| border: 1px solid #DDD; | |
| height: 48px; | |
| min-height: 48px; | |
| } | |
| .stButton > button { | |
| border-radius: 8px; | |
| background-color: #4A90E2; | |
| color: white; | |
| border: none; | |
| padding: 0.75rem 1rem; | |
| font-weight: bold; | |
| width: 100%; | |
| height: 48px; | |
| min-height: 48px; | |
| } | |
| .stButton > button:hover { | |
| background-color: #357ABD; | |
| } | |
| .stFileUploader > div { | |
| padding: 1rem; | |
| border: 2px dashed #4A90E2; | |
| border-radius: 8px; | |
| text-align: center; | |
| margin-bottom: 1rem; | |
| } | |
| .sidebar-logo { | |
| text-align: center; | |
| margin-bottom: 1rem; | |
| } | |
| .input-row { | |
| display: flex; | |
| gap: 0.5rem; | |
| align-items: end; | |
| margin-top: 1rem; | |
| } | |
| .empty-chat { | |
| text-align: center; | |
| color: #888; | |
| font-style: italic; | |
| padding: 1rem; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # ---- Sidebar ---- | |
| with st.sidebar: | |
| st.markdown("<div class='sidebar-logo'><img src='https://cdn-icons-png.flaticon.com/512/4712/4712027.png' width='100'></div>", unsafe_allow_html=True) | |
| st.title("🤖 RAG Assistant") | |
| st.write("• Upload a PDF 📄") | |
| st.write("• Ask context-specific questions 💬") | |
| st.write("• Powered by **LLaMA 3 + Groq** ⚡") | |
| st.markdown("---") | |
| if st.button("🗑️ Clear Chat", use_container_width=True): | |
| cleanup_session_data() | |
| st.session_state.chat_history = [] | |
| st.rerun() | |
| # ---- Title ---- | |
| st.markdown("<div class='title'>📄 Multimodal Chat Assistant</div>", unsafe_allow_html=True) | |
| # ---- Main container ---- | |
| st.markdown("<div class='main-container'>", unsafe_allow_html=True) | |
| # Upload UI | |
| uploaded_file = st.file_uploader("📤 Upload a PDF file", type=["pdf"]) | |
| if uploaded_file: | |
| with st.spinner("📚 Ingesting document..."): | |
| ingest_pdf(uploaded_file.read(), uploaded_file.name) | |
| st.success("✅ Document ingested successfully! You can now ask questions.") | |
| # Initialize chat history | |
| if "chat_history" not in st.session_state: | |
| st.session_state.chat_history = [] | |
| # Chat history display | |
| if st.session_state.chat_history: | |
| st.markdown("<div class='chat-history'>", unsafe_allow_html=True) | |
| for msg in st.session_state.chat_history: | |
| if msg["role"] == "user": | |
| st.markdown(f"<div class='message-user'><b>🧑 You:</b> {msg['content']}</div>", unsafe_allow_html=True) | |
| else: | |
| st.markdown(f"<div class='message-bot'><b>🤖 Assistant:</b> {msg['content']}</div>", unsafe_allow_html=True) | |
| st.markdown("</div>", unsafe_allow_html=True) | |
| # Input & button | |
| col1, col2 = st.columns([5, 1]) | |
| with col1: | |
| prompt = st.text_input("💬 Type your question", placeholder="Ask about your PDF...", label_visibility="collapsed") | |
| with col2: | |
| send_clicked = st.button("Send ➡️", use_container_width=True) | |
| # Response handling | |
| if (prompt and send_clicked) or (prompt and st.session_state.get('last_prompt', '') != prompt): | |
| if prompt.strip(): | |
| st.session_state.chat_history.append({"role": "user", "content": prompt}) | |
| with st.spinner("🤔 Thinking..."): | |
| try: | |
| response = qa_chain.run(prompt) | |
| st.session_state.chat_history.append({"role": "assistant", "content": response}) | |
| except Exception as e: | |
| error_msg = f"Sorry, I encountered an error: {str(e)}" | |
| st.session_state.chat_history.append({"role": "assistant", "content": error_msg}) | |
| st.session_state.last_prompt = prompt | |
| st.rerun() | |
| st.markdown("</div>", unsafe_allow_html=True) | |