Update app.py

app.py

@@ -1,57 +1,209 @@
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import os
-import subprocess
+import streamlit as st
 from dotenv import load_dotenv
+import wikipedia
+from xhtml2pdf import pisa
+import io
+
+# Load environment variables
+load_dotenv()
+
+st.set_page_config(page_title="Ask Wikipedia", page_icon="📘", layout="wide")
+
+# --- Wikipedia Summary ---
+def get_wikipedia_summary(query):
+    try:
+        return wikipedia.summary(query, sentences=2)
+    except wikipedia.exceptions.DisambiguationError as e:
+        return f"Your query is ambiguous, here are some options: {e.options}"
+    except wikipedia.exceptions.HTTPTimeoutError:
+        return "Request timed out. Please try again later."
+    except Exception as e:
+        return f"An error occurred: {e}"
+
+# --- PDF Generation ---
+def generate_pdf(convo, topic):
+    html = f"<h2>{topic}</h2><hr>"
+    for msg in convo:
+        if msg["role"] == "user":
+            html += f"<p><strong>👤 You:</strong> {msg['text']}</p>"
+        elif msg["role"] == "assistant":
+            html += f"<p><strong>📘 Wikipedia:</strong> {msg['text']}</p>"
+
+    result = io.BytesIO()
+    pisa_status = pisa.CreatePDF(io.StringIO(html), dest=result)
+    if pisa_status.err:
+        return None
+    return result
+
+# --- Session Init ---
+if "chat_sessions" not in st.session_state:
+    st.session_state.chat_sessions = {}
+if "current_conversation" not in st.session_state:
+    st.session_state.current_conversation = None
+if "edit_mode" not in st.session_state:
+    st.session_state.edit_mode = {}
+
+# --- App Title ---
+st.title("💬 Ask Wikipedia - Your Knowledge Companion")
+
+# --- Custom CSS ---
+st.markdown("""
+<style>
+.chat-wrapper {
+    display: flex;
+    flex-direction: column;
+}
+
+.message-row {
+    display: flex;
+    align-items: flex-end;
+    justify-content: flex-end;
+    margin-top: 10px;
+}
+
+.user-bubble {
+    background-color: rgba(0, 200, 83, 0.1);
+    color: var(--text-color);
+    padding: 15px;
+    border-radius: 15px;
+    max-width: 80%;
+    border: 1px solid rgba(0, 200, 83, 0.4);
+    position: relative;
+    margin-bottom: 20px;
+}
+
+.assistant-bubble {
+    background-color: rgba(3, 169, 244, 0.1);
+    color: var(--text-color);
+    padding: 15px;
+    border-radius: 15px;
+    max-width: 80%;
+    align-self: flex-start;
+    margin-right: auto;
+    border: 1px solid rgba(3, 169, 244, 0.3);
+    margin-top: 10px;
+    margin-bottom: 40px;
+}
+
+.bubble-header {
+    font-size: 14px;
+    font-weight: bold;
+    margin-bottom: 5px;
+    display: flex;
+    align-items: center;
+}
+
+.bubble-header span {
+    margin-left: 5px;
+}
+
+.icon-col {
+    display: flex;
+    flex-direction: column;
+    gap: 5px;
+    align-items: center;
+}
+</style>
+""", unsafe_allow_html=True)
+
+# --- Sidebar: Conversations ---
+st.sidebar.title("Conversations")
+titles = list(st.session_state.chat_sessions.keys())
+
+if titles:
+    for topic in titles:
+        col1, col2 = st.sidebar.columns([0.8, 0.2])
+        if col1.button(f"🎨 {topic}", key=f"select_{topic}"):
+            st.session_state.current_conversation = topic
+            st.rerun()
+        if col2.button("🗑", key=f"delete_{topic}"):
+            del st.session_state.chat_sessions[topic]
+            if st.session_state.current_conversation == topic:
+                st.session_state.current_conversation = None
+            st.rerun()
+else:
+    st.sidebar.write("No conversations yet. Start one below!")
+
+# --- New Conversation ---
+new_topic = st.sidebar.text_input("New Conversation Name")
+if st.sidebar.button("Start New Conversation"):
+    if new_topic.strip() and new_topic not in st.session_state.chat_sessions:
+        st.session_state.chat_sessions[new_topic] = []
+        st.session_state.current_conversation = new_topic
+        st.sidebar.success(f"Started new conversation: {new_topic}")
+        st.rerun()
+    elif not new_topic.strip():
+        st.sidebar.warning("Please enter a name.")
+    else:
+        st.sidebar.warning("Conversation already exists!")
+
+# --- Main Chat Area ---
+if st.session_state.current_conversation:
+    convo = st.session_state.chat_sessions[st.session_state.current_conversation]
+
+    st.markdown('<div class="chat-wrapper">', unsafe_allow_html=True)
+
+    for idx, msg in enumerate(convo):
+        with st.container():
+            if msg["role"] == "user":
+                if st.session_state.edit_mode.get(idx, False):
+                    new_text = st.text_input("Edit your message:", value=msg["text"], key=f"edit_input_{idx}")
+                    col1, col2 = st.columns([1, 1])
+                    with col1:
+                        if st.button("✅ Save", key=f"save_{idx}"):
+                            msg["text"] = new_text
+                            try:
+                                new_response = get_wikipedia_summary(new_text)
+                            except:
+                                new_response = "Failed to retrieve summary."
+                            if idx + 1 < len(convo) and convo[idx + 1]["role"] == "assistant":
+                                convo[idx + 1]["text"] = new_response
+                            st.session_state.edit_mode[idx] = False
+                            st.rerun()
+                    with col2:
+                        if st.button("❌ Cancel", key=f"cancel_{idx}"):
+                            st.session_state.edit_mode[idx] = False
+                            st.rerun()
+                else:
+                    col1, col2 = st.columns([0.1, 0.9])
+                    with col1:
+                        if st.button("✏️", key=f"edit_btn_{idx}"):
+                            st.session_state.edit_mode[idx] = True
+                            st.rerun()
+                    with col2:
+                        st.markdown(f'''
+                        <div class="user-bubble">
+                            <div class="bubble-header">👤 <span>User</span></div>
+                            {msg["text"]}
+                        </div>
+                        ''', unsafe_allow_html=True)
+
+            elif msg["role"] == "assistant":
+                st.markdown(f'''
+                <div class="assistant-bubble">
+                    <div class="bubble-header">📘 <span>Wikipedia</span></div>
+                    {msg["text"]}
+                </div>
+                ''', unsafe_allow_html=True)
+
+    st.markdown('</div>', unsafe_allow_html=True)
+
+    # --- Export PDF ---
+    if st.button("📥 Export Conversation as PDF"):
+        pdf_bytes = generate_pdf(convo,st.session_state.current_conversation)
+        if pdf_bytes:
+            st.download_button("Download PDF", pdf_bytes, file_name="AskWikipedia_Conversation.pdf", mime="application/pdf")
+        else:
+            st.error("❌ Failed to generate PDF.")
 
-# Function to load model, tokenizer, and dataset only once
-@st.cache_resource  # This decorator caches the loading process
-def load_resources():
-    load_dotenv()
-
-    huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-
-    # Set the environment variable using os
-    # huggingface_token = st.secrets["huggingface"]["token"]
-
-    # Run the huggingface-cli login command from the Python script using subprocess
-    subprocess.run(["huggingface-cli", "login", "--token", huggingface_token])
-    # Load the pre-trained Llama3 model (or your fine-tuned model)
-    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
-    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
-    print("Model and tokenizer loaded.")
-    return model, tokenizer
-
-# Load resources only once and cache them
-if "model" not in st.session_state:
-    with st.spinner("Loading the model... Please wait!"):
-        model, tokenizer = load_resources()
-        st.session_state.model = model
-        st.session_state.tokenizer = tokenizer
-
-# Streamlit UI Setup
-st.title("Llama3-based Chatbot")
-st.write("Ask a question, and I will provide an answer based on the model.")
-
-# User input for the query
-user_query = st.text_input("Your Query:", "")
-
-if user_query:
-    with st.spinner("Generating response..."):
-        # Retrieve the tokenizer from session state
-        tokenizer = st.session_state.tokenizer
-        model = st.session_state.model
-        
-        # Tokenize the input query
-        inputs = tokenizer(user_query, return_tensors="pt")
-
-        # Generate a response from the model
-        with torch.no_grad():
-            outputs = model.generate(inputs["input_ids"], max_length=150, num_return_sequences=1)
-
-        # Decode the response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Display the response
-        st.write(f"Model Response: {response}")
+    # --- User Prompt ---
+    user_input = st.chat_input("Ask Wikipedia...")
+    if user_input:
+        convo.append({"role": "user", "text": user_input})
+        try:
+            reply = get_wikipedia_summary(user_input)
+        except:
+            reply = "Could not fetch response."
+        convo.append({"role": "assistant", "text": reply})
+        st.rerun()